O salto explosivo de desempenho das GPUs
(hazyresearch.stanford.edu)- Com o aumento do custo computacional de IA, a Hazy Research resume que a chave para otimizar o desempenho da GPU é manter os tensor cores da NVIDIA H100 sempre ocupados
- A H100 entrega 989 TFLOPs em multiplicação de matrizes com half-precision, mas fica em cerca de 60 TFLOPs em operações gerais, então a taxa de utilização despenca quando os tensor cores param
- Para chegar perto do desempenho máximo, é preciso tratar em conjunto WGMMA, organização da shared memory, geração de endereços e occupancy; sem
wgmma.mma_async, microbenchmarks ficam em cerca de 63% do pico - O DSL embutido em CUDA ThunderKittens, agora público, usa abstrações de tile e vector para esconder complexidades como swizzling e register layout, simplificando a escrita de kernels da família FlashAttention
- O kernel forward do FlashAttention-2 para H100 é escrito com cerca de 100 linhas, fica cerca de 30% mais rápido que o FlashAttention-2, e o kernel de Based linear attention opera a 215 TFLOPs
Condições que determinam o desempenho da H100
- IA usa muito compute, e a Hazy Research vem trabalhando nos últimos anos para fazer a IA usar menos compute ou rodar com mais eficiência dentro do compute disponível
- Exemplos de redução de compute: Based, Monarch Mixer, H3, Hyena, S4
- Exemplos de execução eficiente: FlashAttention, FlashAttention-2, FlashFFTConv
- O objetivo prático é organizar o que foi aprendido ao tornar GPUs mais rápidas e divulgar o DSL embutido em CUDA ThunderKittens, que ajuda a escrever kernels rápidos
- Em termos mais amplos, o texto aborda como entender o hardware mudou a forma de enxergar o compute de IA
Estrutura e gargalos da NVIDIA H100
- A discussão toma como referência a GPU H100 SXM com a seguinte configuração
- 80GB HBM3, largura de banda de 3TB/s
- 50MB de cache L2, largura de banda de 12TB/s, dividido em duas seções de 25MB conectadas por crossbar em toda a GPU
-
132 SMs
- Cada SM tem 256KB de cache L1, incluindo até 227KB de shared memory, com cerca de 33TB/s de largura de banda combinada
- O novo hardware do Hopper, o Tensor Memory Accelerator (TMA), cuida da geração assíncrona de endereços e do fetch de memória
- Cada SM é composto por 4 quadrants, e cada quadrant tem warp scheduler, 512 vector registers, tensor core para multiplicação de matrizes e instruções embutidas em paralelo
- Todo o compute acontece no SM, e a maior parte é processada em registers
- A chave para extrair desempenho na H100 é manter os tensor cores sempre bem alimentados
- A H100 oferece 989 TFLOPs em multiplicação de matrizes com half-precision e cerca de 60 TFLOPs em operações “não tensor core”
- Nos ciclos em que tensor cores são usados, o hardware atinge pelo menos 94% de utilização
- Nos ciclos em que tensor cores não são usados, a utilização fica em no máximo 6%
WGMMA: instrução necessária, mas difícil de usar
- A H100 traz a instrução de warp group matrix multiply accumulate
wgmma.mma_async- Em PTX:
wgmma.mma_async - Em SASS:
HGMMA/IGMMA/QGMMA/BGMMA
- Em PTX:
- Nas GPUs anteriores,
wmma.mma.syncemma.syncfuncionavam de forma síncrona: um único warp de 32 threads enviava dados ao tensor core e esperava o resultado wgmma.mma_asyncfaz 128 threads contíguas cooperarem e se sincronizarem em todos os quadrants do SM, iniciando diretamente da shared memory uma multiplicação de matrizes assíncrona- Enquanto a multiplicação acontece, os warps podem executar outras tarefas em registers
- O resultado pode ser aguardado no momento desejado
- Em microbenchmarks, essas instruções foram necessárias para extrair todo o compute da H100
- Sem elas, observou-se que a GPU fica em cerca de 63% da utilização de pico
- Isso pode acontecer porque os tensor cores exigem pipelines de hardware profundos até mesmo em recursos locais
- A maior dificuldade é a complexidade do memory layout
- Um layout de shared memory sem swizzle tem coalescing muito ruim e exige muita largura de banda de L2
- O layout com swizzle demorou a ser entendido porque a documentação estava incorreta
- O layout com swizzle parece funcionar apenas para certos formatos de matriz e não combina bem com outros recursos de
wgmma.mma_async - O hardware consegue fazer transpose de submatrizes no caminho até o tensor core, mas só quando o layout não é swizzled
- Em kernels como FlashAttention, TMA e cache L2 são rápidos o suficiente para esconder parte desse problema
- Para usar o hardware por completo, é importante controlar o layout para coalescer memory requests e evitar bank conflicts
Shared memory e bank conflict
- A latência de acesso único da shared memory parece ser de cerca de 30 cycles, tempo em que os tensor cores de um SM quase conseguem executar duas multiplicações de matriz quadrada 32x32
- Trabalhos anteriores como FlashAttention focavam principalmente no gargalo HBM-SRAM, e no passado esse gargalo realmente era importante
- Como a HBM ficou mais rápida e os tensor cores cresceram mais rápido do que outras partes do chip, até a pequena latência da shared memory virou algo a ser eliminado ou escondido
- A shared memory é dividida em 32 banks, então é fácil cair em bank conflict se não houver cuidado
- Quando vários pedaços diferentes de memória no mesmo bank são requisitados ao mesmo tempo, os pedidos são serializados
- Na prática, isso pode deixar o kernel desproporcionalmente mais lento
- O register layout exigido por WGMMA e MMA pode causar bank conflict se usado de forma ingênua
- A solução é reorganizar a shared memory com vários padrões de swizzling para evitar conflitos
- Sempre que possível, é melhor evitar movimentações entre register e shared memory; quando necessárias, vale usar hardware embutido assíncrono como WGMMA e TMA
- Fazer movimentação síncrona com warps reais é o caso mais geral, mas fica perto do pior fallback possível
Geração de endereços e TMA
- Na H100, tensor cores e memória são tão rápidos que a própria geração dos endereços de memória a buscar consome uma parte relevante dos recursos do chip
- Isso fica ainda mais evidente com padrões interleaved complexos ou swizzling
- O Tensor Memory Accelerator (TMA) da NVIDIA permite definir layouts tensoriais multidimensionais em global/shared memory, fazer fetch assíncrono de subtiles desse tensor e disparar uma barrier ao concluir
- O TMA reduz o custo da geração de endereços e também facilita montar pipelines
- Assim como
wgmma.mma_async, o TMA é tratado como essencial para destravar todo o potencial da H100- Pela experiência relatada, ele pode ser até mais importante que o WGMMA
- Poupa recursos de register e instruction dispatch
- Também oferece redução assíncrona em global memory, algo útil em kernels de backward mais complexos
- Entender os modos de swizzling do TMA também exigiu certa engenharia reversa, mas foi menos doloroso do que no WGMMA
Occupancy e os custos que ela esconde
- Em CUDA, occupancy significa a quantidade de threads coescalonadas no mesmo hardware de execução
- O warp scheduler de cada quadrant do SM tenta emitir uma instrução a cada cycle para algum warp pronto para recebê-la
- A H100 depende menos de occupancy do que gerações anteriores
- Graças aos recursos assíncronos, até um único instruction stream pode manter ocupados ao mesmo tempo fetch de memória, multiplicação de matrizes, redução em shared memory e matemática em registers
- Ainda assim, occupancy é muito útil para esconder erros e custos de sincronização
- Um pipeline perfeitamente projetado pode ser rápido mesmo sem occupancy adicional
- Na prática, a observação foi de que as GPUs da NVIDIA parecem projetadas tendo occupancy em mente
- Como sincronização e erros são frequentes, aumentar a occupancy muitas vezes melhora a utilização real do hardware
- Na H100, occupancy ajuda em um nível útil, mas no A100 e na RTX 4090 ela parece ainda mais importante
- A hipótese é que essas GPUs dependam mais de instruction dispatch síncrono do que a H100
ThunderKittens: um pequeno DSL dentro de CUDA
- ThunderKittens é um DSL embutido em CUDA criado para facilitar a escrita de kernels rápidos na H100
- No início ele foi feito para uso interno do laboratório e depois acabou sendo aberto ao público
- O nome foi escolhido porque kittens são fofos e porque acharam divertido escrever
kittens::no código - O ThunderKittens busca simplicidade e oferece quatro tipos com template
- Register tiles: tensor 2D sobre o arquivo de registers
- Register vectors: tensor 1D sobre o arquivo de registers
- Shared tiles: tensor 2D dentro da shared memory
- Shared vectors: tensor 1D dentro da shared memory
- Tile é parametrizado por height, width e layout
- Register vector é parametrizado por length e layout, enquanto shared vector usa apenas length
- Shared vector em geral não sofre com bank conflict
- As operações oferecidas manipulam tiles e vectors no nível de warp ou de grupo cooperativo de warps
- initializer: como zerar um shared vector
- unary op: como
exp - binary op: como
mul - row/column op: como
row_sum
- Como o ThunderKittens é embutido em CUDA, ao contrário de bibliotecas como Triton, a abstração pode falhar de forma “graceful” segundo a descrição
- Se faltar alguma funcionalidade, dá para estender do jeito desejado
Exemplo com FlashAttention e desempenho
- Como exemplo do ThunderKittens, é mostrado um kernel forward simples de FlashAttention para RTX 4090
- Lida apenas com headdim=64
nprecisa ser múltiplo de 256- É escrito em cerca de 60 linhas de código CUDA
- A utilização do hardware é de 75%
- A maior parte da complexidade está no próprio algoritmo, não em swizzling ou register layout
- O forward pass do FlashAttention-2 para H100 também foi escrito com ThunderKittens
- O ThunderKittens encapsula a complexidade de TMA, WGMMA, modos de swizzling e descriptors
- O kernel tem cerca de 100 linhas
- Na H100, ele fica cerca de 30% mais rápido que o FlashAttention-2
- O ThunderKittens é descrito como um “mini-pytorch” para GPU, encapsulando layouts e instruções e oferecendo primitives
- Também foram divulgados o kernel de Based linear attention e kernels para outras arquiteturas que serão publicados em breve
- O kernel de Based linear attention roda a 215 TFLOPs
- Considerando o recompute do próprio algoritmo, passa de 300 TFLOPs
- Linear attention é teoricamente mais eficiente, mas historicamente teve eficiência real muito menor no hardware
- O resultado pode ampliar o conjunto de aplicações de alto throughput
Pensar em torno de tiles
- A visão apresentada é que o ThunderKittens funciona bem justamente por não tentar fazer tudo
- CUDA é muito mais expressivo do que o ThunderKittens
- O ThunderKittens é um DSL pequeno e simples
- A abstração central é o small tile, algo considerado alinhado à direção do hardware e da IA
- O ThunderKittens não oferece suporte a dimensões menores que 16
- A visão é que o hardware também não quer particularmente dimensões tão pequenas
- A provocação é algo como: “se sua multiplicação de matrizes é menor que 16x16, dá mesmo para ter certeza de que isso é IA?”
- A visão herdada da era da CPU, em que uma word de 32 bits é o register, não parece adequada para hardware de IA
- O vector register de 1024 bits do CUDA é visto como um passo na direção certa
- Nesse contexto, o register passa a ser o dado de um tile 16x16
- IA continua girando em torno de multiplicação de matrizes, reduction e reshape, então a abstração de tile se encaixa tanto na IA quanto no hardware
- Daqui para frente, será preciso reorganizar ideias de IA de formas que se mapeiem bem ao hardware
- O tamanho do estado recorrente precisa ser grande o bastante para caber em um SM
- A densidade de compute não pode ficar abaixo do nível exigido pelo hardware
- Ajustar o que foi aprendido com o hardware ao design de IA é apontado como uma direção importante
Planos de suporte à AMD
- O suporte do ThunderKittens a hardware AMD deve chegar em breve
1 comentários
Comentários do Hacker News
A pergunta "se a multiplicação de matrizes é menor que 16x16, isso é mesmo IA?" é interessante
Os requisitos de hardware para IA estão ficando cada vez mais claros. A GPU foi originalmente projetada para um propósito totalmente diferente, mas acabou sendo usada em IA porque o hardware de multiplicação de matrizes era bom, e uma "GPU de IA" pode remover algumas funções presentes em uma GPU real
Também há uma tendência de representações numéricas mais curtas, como ponto flutuante de 16 bits, 8 bits, 2 bits e 1 bit, e em algum momento um ponto ideal será definido. Este texto mostra que um hardware que prefere blocos 16x16 faz bastante sentido. Alguém já deve estar escrevendo isso em VHDL ou provavelmente fará isso em breve
No fim, parece provável que surja um dispositivo mais simples, menos genérico e mais barato, capaz de executar apenas operações de "IA" com o mínimo possível de peso desnecessário de hardware
A Nvidia provavelmente também está trabalhando nisso, mas do ponto de vista de negócios talvez seja melhor manter a forma de placa de vídeo, um dispositivo que reúne jogos/entretenimento/criptomoedas/IA
[1] https://github.com/hollance/neural-engine/blob/master/docs/a...
Isso lembra a época em que a Nervana de Naveen Rao fazia um driver Nvidia Maxwell mais rápido que o próprio driver da Nvidia. Nem todo erro de documentação em um produto que cresce rápido é uma resposta à concorrência, mas considerando que os pesquisadores demoraram bastante para fazer engenharia reversa do wgmma e até a situação política entre EUA e China em torno do H100, parece que a Nvidia está usando seus velhos truques para proteger seu fosso competitivo
Por isso, em vez de se aprofundar demais nas peculiaridades do H100, é preciso ver que "qual é o hardware que a IA quer" também inclui a situação comercial
https://www.amd.com/en/products/accelerators/alveo/v80.html
Arquitetura XDNA
https://www.amd.com/en/technologies/xdna.html
O trecho "Mentiras da NVIDIA. Isso é uma descrição tremendamente enganosa do layout real de wgmma com swizzle de 128b. Estou expondo publicamente porque esse diagrama me custou três semanas irrecuperáveis da minha vida" chama atenção
Fico imaginando se alguém se surpreenderia ao saber que uma parte enorme dos avanços em IA está em engenharia como otimização de multiplicação de matrizes, e que uma parte considerável dessa engenharia é engenharia reversa de chips da NVIDIA
Escalonador de warps, quatro quadrantes, acelerador de memória tensor, layout wgmma sem swizzle…
A fronteira entre a terminologia de GPU e o tecnoblablá ao estilo Star Trek está ficando cada vez mais tênue
Às vezes penso isso ao ver outros textos também. Fico imaginando que impressão teria alguém que recebesse um link para um texto daqui e fosse ler. Pareceria entrar num evento de fãs de Trek discutindo o núcleo de dobra
Para reduzir o consumo de energia da inferência de IA e aumentar a velocidade, parece que o melhor seria migrar para circuitos analógicos aproximados
Não é necessário multiplicação e soma em ponto flutuante perfeitas; basta um dispositivo que receba duas tensões de entrada e produza uma tensão de saída suficientemente próxima do resultado da multiplicação
A grande vantagem é representar o número por meio da tensão em um único fio, em vez de usar 16 fios para representar um float16. Em teoria, talvez seja possível até uma precisão muito maior que float32. Além disso, como os valores podem ser conectados diretamente sem serem carregados em uma unidade lógico-aritmética, a economia potencial de área de die e de energia pode chegar a várias ordens de grandeza
Por exemplo, aceitar que um em cada milhão de bits de saída inverta e, em troca, melhorar a relação desempenho/energia. Seria difícil com float32, em que um único valor infinito pode estragar tudo, mas em int8 talvez dê para tolerar que às vezes saia 128 quando se queria 0
[1] Não sei se a unidade de ponto flutuante matricial da H100 realmente segue o padrão IEEE 754
Redes neurais biológicas não são quase totalmente conectadas como redes neurais artificiais comuns, e os coeficientes de conexão de entrada e saída dos neurônios são muito locais, com menos de 10. Pelo que sabemos, na biologia também não existe retropropagação; em vez disso, há feedback e recorrência
Pode haver células auxiliares ou processos ainda desconhecidos que sejam essenciais para o funcionamento do sistema nervoso central. Mesmo em alto nível, talvez exista uma quantidade considerável de conectividade "hardcoded", e parte disso já é conhecida. Por exemplo, os neurônios auditivos do ouvido são conectados, e acontece algo parecido com uma convolução para localizar a posição do som. Isso não é um fenômeno emergente, e sim uma capacidade possível mesmo sem treinamento
A vida descobriu isso ao longo de bilhões de anos e um número semelhante de gerações, então não é surpreendente. Em teoria isso também poderia ser feito em software, mas, considerando mais de 1 trilhão de neurônios no cérebro de primatas/humanos, é extremamente difícil até com máquinas atuais na faixa de mil núcleos. Mesmo na "nuvem", não seria possível satisfazer a conectividade e a latência necessárias
Seria legal se uma abordagem dessas conseguisse modelar com sucesso algo no nível de um verme ou inseto
Este texto me fez lembrar da diversão que senti na aula de programação paralela de CS 149
O estilo deste texto é realmente impressionante, e estou ansioso para ver isso em AMD MI300x. Se quiser gastar um tempo com o meu equipamento, é só me avisar
Queria saber quão bem isso realmente funciona, ou se vale mais a pena juntar um pouco mais e comprar a XTX em vez da 7900 XT, e o quanto ter menos VRAM afeta a usabilidade real
O leitor não deveria ter que ir até o knowyourmeme.com para entender o que os autores querem dizer. Eu nem sei o que este título quer dizer, e acho que por isso ele erra feio o alvo
Fico curioso sobre por onde começar e qual roadmap seguir para entender totalmente um texto como este
E vale a pena tentar escrever você mesmo um kernel CUDA que faça multiplicação vetor-matriz. Com pycuda, você pode focar no kernel e escrever o resto em Python. Basta dizer ao ChatGPT que quer criar por conta própria uma implementação que multiplique um vetor de 4000 elementos por uma matriz 4000x12000 e pedir que ele guie todo o processo
Para alugar GPUs, Runpod é bom, e hoje tem desde GPUs baratas até H100. No começo, dá para começar com uma GPU de nível mais baixo
Passei 2 meses implementando e otimizando kernels de multiplicação de matrizes com Spiral
O gráfico no README do GitHub (https://github.com/HazyResearch/ThunderKittens/blob/main/att...) está bagunçado demais. Esses gráficos de barras onduladas são mesmo legais? :P
[1]: https://matplotlib.org/stable/gallery/showcase/xkcd.html#sph...
O nome ThunderKittens é excelente. Quero ver o ThunderKittens lidando com o backward do FlashAttention, que é uma ordem de grandeza mais difícil do que o forward
causal: https://github.com/HazyResearch/ThunderKittens/blob/main/exa...
non-causal: https://github.com/HazyResearch/ThunderKittens/blob/main/exa...
Esse tipo de pesquisa não é algo que as equipes que fazem NPUs já vêm fazendo hoje em dia? Por exemplo, os chips da Groq conseguem o desempenho atual porque usam uma arquitetura dedicada para IA. No lado do consumidor, o Apple Silicon também é bastante competente
Não sou da área, mas parece que há um limite para o que dá para fazer só com processadores de propósito geral que se comunicam por caminhos relativamente lentos. Repensar o projeto no nível do hardware e, no fim, reduzir o preço no mercado consumidor parece uma estratégia melhor no longo prazo
Dá para comprar uma GPU da Nvidia por algumas centenas de dólares, ou um notebook gamer com 4050 e 6 GB de VRAM por 900 dólares, então é difícil chamar IA baseada em CPU de competente
No trabalho também tentei usar CPU por falta de GPU, mas fora usar modelos pequenos e esperar, não era algo realista. No fim, acabei pedindo um computador com GPU
"Tecnicamente possível" e "realmente bom de usar" são coisas diferentes. A Nvidia foi realmente boa de usar, e CPU foi doloroso e frustrante