O avanço explosivo de desempenho das GPUs

Características da GPU H100

• Oferece 80GB de memória HBM3 e largura de banda de 3TB/s (na prática, é um pouco menor)
• Oferece 50MB de cache L2 e largura de banda de 12TB/s. É dividido na GPU em duas seções de 25MB conectadas por um crossbar (o crossbar é um fator de perda de desempenho)
• É composta por 132 Streaming Multiprocessors (SM), e cada SM tem a configuração abaixo:
  • Oferece até 227KB de memória compartilhada dentro de 256KB de cache L1 (juntos, cerca de 33TB/s de largura de banda)
  • Oferece Tensor Memory Accelerator (TMA), que permite geração assíncrona de endereços e fetching de memória. Também fornece recursos como suporte à rede de memória on-chip, mas isso não será abordado neste post
  • É dividido em 4 quadrantes, e cada quadrante é composto por warp scheduler, 512 vector registers (cada um contendo 32 palavras de 4 bytes), tensor core para multiplicação de matrizes e instruções embutidas que suportam operações paralelas como sum/multiply

Dicas de otimização de desempenho da GPU H100

• É importante aproveitar ao máximo os Tensor Cores. A H100 tem desempenho de 989 TFLOPs em multiplicação de matrizes FP16, portanto o nível de uso dos Tensor Cores afeta fortemente a taxa de utilização da GPU como um todo
• Porém, para aproveitar ao máximo os Tensor Cores, é preciso considerar os pontos abaixo:
  • O uso da instrução Warp Group Matrix Multiply Accumulate (WGMMA) é essencial, mas difícil de usar
  • A Shared Memory não é tão rápida quanto parece, e seu uso exige bastante cuidado
  • A geração de endereços tem custo alto, então é preciso otimizar usando recursos como Tensor Memory Accelerator (TMA)
  • Aumentar a occupancy continua ajudando, e os registers são o principal recurso
• Essas características se aplicam até certo ponto não apenas à H100, mas também a outras GPUs; ainda assim, na H100 o uso dos Tensor Cores é especialmente importante e desafiador

ThunderKittens: DSL embarcada de CUDA otimizada para a H100

• DSL embarcada baseada em CUDA desenvolvida para extrair o máximo desempenho de GPUs modernas como a NVIDIA H100
• Fornece os 4 tipos de template baseados em tile abaixo:
  • Register Tile (tensor 2D armazenado em registradores)
  • Register Vector (tensor 1D armazenado em registradores)
  • Shared Tile (tensor 2D armazenado em Shared Memory)
  • Shared Vector (tensor 1D armazenado em Shared Memory)
• Também fornece vários operadores para manipulação de tiles (exp, mul, sum etc.) no nível de warp ou de grupo de warps
• Ao implementar os kernels de Flash Attention e Flash Attention 2 com ThunderKittens, o código fica muito mais conciso e o desempenho na H100 melhora em até 30%
• O kernel de Based Linear Attention também foi implementado com ThunderKittens, alcançando desempenho de 215 TFLOPs (incluindo recompute, por características do algoritmo, passa de 300 TFLOPs)

Reflexão sob uma perspectiva filosófica

• O motivo de o ThunderKittens funcionar bem é que ele não tenta dar suporte a tudo, e sim fornece uma abstração baseada em tiles que é simples e se encaixa bem na arquitetura da GPU
• Usar vector registers de 1024 bits no lugar das tradicionais palavras de 32 bits também é um avanço, mas é necessária uma mudança de paradigma para enxergar tiles 16x16 como a unidade dos registradores
• Como cargas de trabalho de IA acabam se concentrando em multiplicação de matrizes, reduction e reshape, a abordagem baseada em tiles faz sentido, e do ponto de vista de hardware a evolução deve seguir não apenas com systolic arrays, mas também com suporte a pequenas multiplicações de matrizes
• Mais adiante, é preciso mudar a forma de pensar para projetar algoritmos de IA em formatos otimizados para hardware. Por exemplo, limitar o tamanho do estado de uma RNN ao que cabe em um SM e ajustar também a densidade computacional ao nível exigido pelo hardware

Opinião do GN⁺

• O ThunderKittens pode ser uma opção atraente para desenvolvedores já familiarizados com CUDA. Isso porque é possível obter ganhos de desempenho com facilidade sem mexer muito no código existente dos kernels
• Ainda assim, para iniciantes a barreira de entrada pode continuar alta. Parece que será necessário oferecer mais exemplos de código e materiais de aprendizado no futuro
• É interessante o plano de expandir o suporte do ThunderKittens não só para a H100, mas também para outras GPUs como as da AMD. Isso pode ajudar a reduzir a dependência de fornecedor
• Em última instância, um ponto muito importante é projetar o próprio modelo/algoritmo de IA em um formato otimizado para hardware. Para isso, é preciso antes ter uma compreensão profunda das características do hardware, e o ThunderKittens pode ajudar desenvolvedores a obter esse tipo de insight
• Espera-se que a pesquisa e o desenvolvimento contínuos da Hazy Research, bem como suas contribuições open source, ajudem bastante a fortalecer o ecossistema CUDA. Ainda assim, no longo prazo, também parece necessária a chegada de frameworks com nível de abstração mais alto

[GN⁺]

Comentários no Hacker News

• Os requisitos de hardware para IA estão ficando cada vez mais claros. As GPUs foram originalmente projetadas para outros propósitos, mas estão sendo usadas em IA por terem um bom hardware para multiplicação de matrizes. Uma "GPU de IA" poderia excluir alguns recursos de uma GPU real, e há uma tendência em direção a números menores (ponto flutuante de 16 bits, 8 bits, 2 bits e 1 bit). Este artigo sugere que hardware que prefere blocos 16x16 tem muitas vantagens.

• Há necessidade de coprocessadores dedicados para IA (NPUs), especialmente em sistemas desktop prosumer para desenvolvedores, especialistas, gamers etc. As GPUs funcionam no ambiente corporativo, mas são inconvenientes para uso de IA na computação pessoal. Em especial, as limitações de VRAM e a ausência de APIs abertas padronizadas além de Vulkan são um problema.

• Para avançar a pesquisa em IA, é preciso estudar melhor neurociência e psicologia. Além disso, questões relacionadas à topologia de grafos das redes neurais também podem ser relevantes.

• É o CUTLASS tornado amigável para o usuário?

• O mascote do ThunderKittens tem uma vibe de gato/Sony Aibo. Parece ter sido bem gerado por IA.

• Se a computação básica universal (UBC) fosse considerada um substituto para a renda básica universal, isso seria um futuro bastante distópico. Imagine uma única empresa como a Nvidia produzindo toda a computação.