Como fazer a GPU da AMD rodar na velocidade de ‘brrr’

• HipKittens é um conjunto de primitivas de programação projetado para extrair o desempenho potencial das GPUs AMD, otimizando acesso à memória, escalonamento e reutilização de cache
• A GPU AMD MI355X tem uma estrutura com 256 compute units e 8 chiplets (XCD), oferecendo um grande arquivo de registradores e instruções de núcleo matricial granulares
• Ao contrário da NVIDIA, a AMD não tem realocação de registradores, instruções matriciais assíncronas nem mbarrier, então, em vez de wave specialization, funcionam melhor os escalonamentos 8-wave ping-pong e 4-wave interleave
• O HipKittens melhora a localidade de cache L2 e LLC com escalonamento de grid consciente de chiplet, alcançando ganhos de largura de banda e TFLOPS em operações de GEMM e Attention
• Essa abordagem compensa a falta de maturidade do software no ecossistema de GPUs AMD e fornece uma base para ampliar a escalabilidade da computação de IA em diferentes hardwares

---

Estrutura e características de desempenho das GPUs AMD CDNA

• A GPU AMD MI355X inclui 256 compute units (CU), e cada CU é composta por 4 SIMD
  • Um SIMD executa uma wave de 64 threads, em contraste com o warp de 32 threads da NVIDIA
• A MI355X tem 70% do SRAM da B200 (165KB) e não possui instruções assíncronas de multiplicação matricial, realocação de registradores, aceleração de memória tensor nem mbarrier
• Em compensação, oferece um arquivo de registradores 2x maior e 60% mais processadores (256 CU vs. 160 SM)
  • Suporta instruções de núcleo matricial pequenas e granulares e tem carregamento direto de memória global para memória compartilhada (semelhante a TMA)
• A AMD adota uma arquitetura de chiplets composta por 8 chiplets (XCD), cada um com cache L2 independente, além de um cache LLC em nível superior
• Segundo a tabela, a MI355X entrega BF16 2,5 PFLOPs, MXFP8 5,0 PFLOPs e MXFP6 10,1 PFLOPs, além de 288GB de memória e 8TB/s de largura de banda

Desafios de projeto de kernels para AMD

• Otimização de acesso à memória: devido às limitações do compilador HIPCC e ao comportamento de I/O não público, o projeto de layout de dados e padrões de swizzle é importante
• Escalonamento dentro do processador: na AMD, é preciso aproveitar o arquivo de registradores e instruções matriciais pequenas em vez de memória compartilhada
• Escalonamento entre processadores: por causa da estrutura baseada em chiplets, é necessário distribuir o trabalho considerando efeitos NUMA no nível de cache

Padrões de acesso à memória do HipKittens

• O HipKittens (HK) usa tiles como unidade básica de dados e fornece funções de operação semelhantes às do PyTorch
  • Um tile é definido por tipo de dado, tamanho e layout, e atende diferentes entradas com metaprogramação de templates em C++
• Escalonamento de registradores: como o HIPCC não consegue usar certos registradores como entrada de MFMA, o HK oferece fixação explícita de registradores
  • O desenvolvedor pode especificar diretamente os registradores para escrever kernels de desempenho máximo
• Layout de registradores: na AMD, o layout muda conforme o tipo de dado e a forma da matriz, então um único padrão de swizzle não é possível
  • Por exemplo, um tile bf16 16×16 e um tile bf16 16×32 exigem padrões de swizzle diferentes
• Estrutura de fases das instruções: as instruções de memória compartilhada da AMD têm grupos de fases descontínuos e pouca documentação interna
  • O HK fornece um solver obtido por engenharia reversa para isso
• Geração de endereços: a AMD suporta carregamento assíncrono de HBM para memória compartilhada e faz otimização com swizzle de endereços HBM

Escalonamento dentro do processador: padrões de wave

• Wave specialization funciona bem na NVIDIA, mas na AMD há perda de desempenho por causa da ausência de realocação de registradores
  • Waves produtoras ocupam registradores desnecessários, enquanto waves consumidoras sofrem spill por falta de registradores
• Nos experimentos do HK, a wave specialization na AMD causou queda de intensidade aritmética e gargalo de memória
  • Exemplo: em GEMM, a configuração HK 0/8 alcançou 1605 TFLOPs, enquanto o CUTLASS chegou a 1570 TFLOPs
• Padrões alternativos de escalonamento
  • 8-wave ping-pong: duas waves executam alternadamente clusters de memória/cálculo
  • 4-wave interleave: uma wave intercala de forma fina operações de memória e cálculo
  • O 8-wave tem código mais simples, enquanto o 4-wave é mais granular, mas gera código mais longo
  • Em GEMM e Attention Forward, o 8-wave atingiu desempenho em nível SoTA

Escalonamento entre processadores: abordagem consciente de chiplet

• A AMD MI355X tem 8 chiplets XCD, e cada chiplet possui cache L2 independente
  • Como os blocos de thread são atribuídos aos chiplets em esquema round-robin, a ordem do grid afeta diretamente a eficiência de reutilização de cache
• Um arranjo simples em row-major tem baixa taxa de reutilização de cache L2, causando perda de largura de banda
  • Exemplo: L2 55%, LLC 95%, 15,1 TB/s, 1113 TFLOPs
• O HK introduz escalonamento de grid consciente de chiplet, aproveitando ao mesmo tempo a localidade dos caches L2 e LLC
  • Ele agrupa blocos de thread por regiões adjacentes da matriz de saída para maximizar a reutilização dos dados de entrada

Exemplos de kernels reais

• Os hot loops dos kernels de Attention Forward e BF16 GEMM usam o esquema 8-wave ping-pong do HK
  • Cada loop executa alternadamente clusters de Compute–Memory e sincroniza com barreiras de escalonamento
  • No exemplo de código, operações do HK como mma_AtB, load, exp2, col_sum são usadas repetidamente

Conclusão: AMD na era da IA multi-silicon

• O HipKittens alcança desempenho competitivo em AMD CDNA3 e CDNA4
  • Três pontos principais: acesso à memória otimizado, escalonamento de wave centrado na AMD e escalonamento de grid consciente de chiplet
• Os kernels do HK alcançam o melhor desempenho no ecossistema AMD e também competem com kernels da NVIDIA Blackwell
• Para ampliar a diversidade da computação de IA, é preciso expandir a acessibilidade às GPUs AMD, e o HipKittens fornece a base de software essencial para isso
• Melhorias no escalonamento de registradores do HIPCC são apontadas como uma área importante de evolução futura

[GN⁺]

Comentários no Hacker News

• Recomendo consultar a discussão sobre HipKittens
• Há também um post sobre a mesma pesquisa, HipKittens: Fast and furious AMD kernels. Há comentários de George Hotz e de funcionários da AMD
• É bom ver a academia lidando com esse tipo de problema, mas acho que, no fim das contas, isso é um problema que a AMD precisa resolver internamente
  • Eu acho melhor que empresas de hardware façam só hardware. Assim, os incentivos permanecem puros. Mesmo que o desempenho caia 20%, ainda acho isso melhor
  • Concordo totalmente. A AMD adiou esse problema por 10 anos e só agora está tentando correr atrás. O hardware é excelente, mas ela não consegue aproveitar seu potencial por falta de capacidade para escrever firmware
  • Mas esse grupo de pesquisa também já criou software semelhante para GPUs da Nvidia. Parece que pesquisadores excelentes estão aplicando sua especialidade
  • Pelo que eu sei, a AMD já está lidando com esse problema em vários níveis e também está colaborando com a tinycorp
• Lendo o texto, fica a impressão de que é difícil otimizar por causa da complexidade arquitetural das GPUs da AMD. Mas, no longo prazo, a abordagem da AMD pode escalar melhor. Enquanto a Nvidia usa 2 chiplets, a AMD tem uma estrutura com 8 chiplets, então há problemas de localidade de memória. Como no futuro o número de chiplets deve aumentar ainda mais, a experiência de lidar com essa complexidade agora pode acabar ajudando no longo prazo
  • A AMD não precisa de warp specialization para obter alto desempenho, então a programação é mais simples
• Muitos desenvolvedores tentaram fazer as GPUs da AMD “go brrr” para o desenvolvedor comum, mas fracassaram. Não entendo por que a AMD não resolve os problemas de software por conta própria. Hoje ela já tem dinheiro suficiente, então não contratar desenvolvedores não é desculpa. As GPUs de datacenter também não são ruins, mas, para quem faz experimentos pessoais de ML e IA, a Nvidia ainda é muito melhor. Sinto que minha RTX 3090 de 5 anos atrás ainda é melhor do que qualquer GPU de consumo da AMD lançada até agora
  • A experiência de desenvolvedor na AMD é terrível. Eles nem aceitam relatórios de bug de crash de driver
  • Recentemente troquei um servidor de inferência com uma NVidia 5090 por duas AMD R9700 32GB, e a experiência foi totalmente positiva. Funcionou de primeira no kernel do Fedora, sem configuração de DKMS, e também foi fácil conectar os contêineres com ROCm. Bastou mudar as configurações do Ollama e do Storyteller. Foi uma experiência muito mais agradável do que com CUDA
  • A Nvidia chega ao ponto de manter até um fork do Unreal Engine. A AMD nem consegue competir nisso
  • A Nvidia é a única empresa de hardware que oferece remuneração competitiva para engenheiros de software. A AMD ainda mantém uma cultura em que software não é visto como “trabalho de verdade”, e esse tipo de inércia é difícil de mudar
• O Mojo tinha uma ideia de melhorar a experiência do desenvolvedor (devX) nas GPUs da AMD, e tenho curiosidade sobre como isso está andando
• Não consigo entender por que a AMD não investe bilhões de dólares em melhorar o software. A Nvidia é a empresa mais valiosa do mundo, e a AMD é a única concorrente dela
  • A AMD também está tentando, mas acho difícil transformar uma cultura organizacional de atualização anual de hardware em uma cultura centrada em software. Software não gera receita imediata como hardware, então a diretoria tende a lhe dar menos prioridade. Além disso, fornecedores externos entregam código em open source, o que parece bom no curto prazo, mas prejudica a qualidade no longo prazo. Se você perder uma única tendência de hardware, corre o risco de ficar para trás em relação aos concorrentes
  • Trabalhei com vários fornecedores de GPU, e só a Nvidia vê software como um ativo (asset) e investe nisso. As outras empresas tratam isso apenas como custo
• Pessoalmente, não gosto muito do meme “go brr”, mas é engraçado ver isso sendo usado em lugares como Stanford
  • Na verdade, eles já tinham usado “go brr” há um ano, no anúncio do ThunderKittens
  • Se esse tipo de meme apareceu num canal oficial de universidade, talvez isso já seja um sinal de que a moda passou
• O projeto em si é excelente, mas fico me perguntando por que a AMD não faz isso diretamente. Parece que a AMD ainda não entende a importância de uma stack de software madura. É preciso uma stack unificada como o CUDA, que funcione em todas as placas. Houve uma época em que eu acreditava que a AMD acabaria alcançando isso, mas agora estou quase desistindo dessa ideia
• O projeto é bom, mas o texto em si passa uma sensação de ter sido escrito de forma estranha
  • O texto é esquisito demais. Parece que houve dependência excessiva de IA, ou uma tentativa de imitar um estilo de escrita de IA. Frases como “veja a part one” ou “como fazer GPUs da AMD go brr” ficam se repetindo. O mais decepcionante é que, nas partes técnicas, algo que deveria ser explicado com gráficos foi descrito em 100 linhas de código