Criando do zero um motor de inferência para LLM com C++ e CUDA

• Como construir um motor de inferência para LLM com C++ e CUDA sem usar bibliotecas
• Com isso, é possível entender toda a stack de inferência de LLM e perceber na prática como diferentes otimizações afetam a velocidade de inferência
• Objetivo: implementar o modelo para inferência rápida em lote único em um servidor com uma única CPU + GPU e atingir uma velocidade de processamento de tokens superior à do llama.cpp

1. Visão geral da arquitetura de LLM e da inferência

• A maioria dos principais LLMs segue a mesma arquitetura, usando blocos de transformer em sequência.
• O carregamento do modelo consiste em definir uma classe de bloco transformer customizável, organizá-la em sequência e inicializá-la com pesos em safetensors.
• A inferência ocorre principalmente em lote único, e a "fase de decodificação" responde pela maior parte da execução.

1.1 Visão geral da inferência

• A inferência se divide na fase de prefill, que envia os tokens do prompt ao modelo para preencher o cache KV, e na fase de decode, que passa repetidamente pelo modelo para gerar tokens
  • Fase de prefill: processa os tokens do prompt e inicializa o cache KV
  • Fase de decode: gera um token por vez
• Cache KV: armazena pares anteriores de chave/valor para calcular rapidamente a atenção com o contexto passado
• O forward pass do modelo usa a tabela de embeddings para mapear IDs de tokens para vetores de embedding e transforma o estado por meio de uma sequência de blocos transformer

1.2 Gargalos e benchmarks

• Gargalo: em hardware moderno, a largura de banda de memória é o fator limitante
  • Ao gerar cada token na inferência do modelo, é necessário ler o modelo inteiro, e a largura de banda de memória é uma limitação maior do que a computação
• A quantização do modelo é eficaz para melhorar a velocidade de inferência
• O throughput máximo teórico de tokens varia conforme o hardware, e o desempenho real pode ser verificado por meio de vários motores de inferência
• Limite teórico de velocidade:
  • AMD EPYC 7702P: máximo de 13.6 tok/s (base FP16)
  • RTX 4090: máximo de 67.1 tok/s (base FP16)
• Benchmarks:
  • llama.cpp: CPU 8.7 tok/s, GPU 61 tok/s
  • calm: GPU 66 tok/s

2. Inferência baseada em CPU

• A implementação inicial na CPU é single-thread e suporta apenas pesos em FP32
• É possível começar a paralelizar o código com multithreading e melhorar o desempenho com SIMD

2.1 Multithreading

• Uso de OpenMP para paralelizar multiplicação matriz-vetor (matmul) e atenção multi-head, melhorando o desempenho
• Resultado da otimização: velocidade melhorou de 0.6 tok/s → 4.4 tok/s

2.2 Quantização de pesos e otimização com SIMD

• Quantização: quantizar pesos FP32 para FP16 reduz pela metade o uso de memória e melhora o desempenho
• SIMD: otimização com AVX2 para processar 8 valores FP32 ao mesmo tempo
• Resultado: 8.4 tok/s alcançados

3. Inferência baseada em GPU

• Ao quantizar o modelo para FP16 e carregá-lo em uma RTX 4090, é possível começar a implementar a inferência em GPU
• Com CUDA, funções em C++ (kernels) podem ser executadas em paralelo na GPU

3.1 Port simples para CUDA

• É possível implementar o backend de GPU convertendo as operações da CPU 1 para 1 em kernels CUDA
• Kernels CUDA são executados de forma assíncrona, mas em uma mesma stream são executados sequencialmente
• Problema: a ineficiência das threads impede o uso adequado dos recursos da GPU → lento, com 2.9 tok/s

3.2 Melhor multiplicação de matrizes (matmul)

• A multiplicação de matrizes ocupa grande parte do tempo de execução na CPU e pode ser otimizada com OpenMP
• Na GPU, é possível aumentar o aproveitamento das threads fazendo cada bloco processar uma linha
• Método de otimização:
  1. Um bloco processa uma linha, e as threads do bloco colaboram no cálculo
  2. Aplicação de warp reduction
• Resultado: velocidade melhorada para 51.7 tok/s

3.3 Fusão de kernels e otimizações adicionais

• É possível melhorar o desempenho fundindo kernels
  • Fusão de kernels: combinar operações consecutivas em um único kernel para minimizar acesso à memória e tempo de computação
• Com otimização do padrão de acesso à memória e reutilização de espaço, alcançou-se 56.1 tok/s

3.4 Otimização de attention e processamento de contexto longo

• Problema: em contextos longos, o kernel de attention vira gargalo de desempenho
• Solução:
  1. Otimização de acesso à memória: redesenhar para ler blocos contíguos de memória
  2. Uso de memória compartilhada em vez de atomicAdd para resolver problemas com valores de ponto flutuante perdidos
• Resultado da otimização:
  • Contexto curto: 63.8 tok/s (mais rápido que os 61.0 tok/s do llama.cpp)
  • Contexto longo: 58.8 tok/s

3.5 Quantização do cache KV e problemas de otimização do compilador

• Quantizar o cache KV para FP16 causa queda de desempenho (falta de otimização do compilador)
• Solução: desenrolar loops manualmente e aplicar prefetch de memória
• Resultado: aproximadamente 2x mais rápido em relação ao FP32 e manutenção de 58.8 tok/s em contexto longo

4. Próximas melhorias possíveis

• Otimização do prefill do prompt: processar vários tokens ao mesmo tempo para reduzir o tempo até o primeiro token gerado
• Fusão do kernel de attention: aplicar técnicas de otimização como FlashAttention
• Quantização mais agressiva: aplicar FP8, INT8, INT4 e quantização de ativação/cache
• Otimização de kernels: introduzir técnicas avançadas para maximizar a largura de banda de memória e a eficiência computacional
• Uso de bibliotecas: aproveitar bibliotecas como cuDNN e cuBLAS para reduzir o tempo de otimização

Resumo dos resultados:

• Velocidade de 63.8 tok/s alcançada com várias otimizações em CPU e GPU
• Desempenho próximo ou superior ao de llama.cpp e calm
• Implementação de um motor de inferência para LLM de alto desempenho usando apenas C++ e CUDA, sem bibliotecas

[GN⁺]

Comentários no Hacker News

• O autor está feliz que a postagem do seu blog tenha chamado atenção e gostaria de ouvir feedback
• Um leitor elogia o texto, dizendo que está excelente, e pergunta quanto tempo levou para escrevê-lo
  • Como alguém que trabalha na área de GPGPU, ele gostaria de escrever algo parecido, mas hesita por causa da incerteza sobre o tempo necessário
• Outro leitor acha que o código não utiliza tensor cores nem instruções wgmma
  • Explica que esse tipo de programação é difícil porque exige lidar com várias tarefas ao mesmo tempo
  • Menciona que, por limitações de largura de banda, talvez não sejam necessários cálculos adicionais
  • Avalia que o código do blog provavelmente funcionará bem ao ser portado para outros aceleradores
  • Demonstra preocupação de que usar wgmma possa prejudicar a portabilidade entre gerações da Nvidia
• Outro leitor está procurando material semelhante em Python para compartilhar com a equipe
  • Quer um material mais completo conceitualmente e conciso no estilo tutorial, em vez de focado em desempenho
• Um usuário quer comparar sua versão do Mistral e o desempenho em tokens por segundo
  • Recomenda-se consultar a seção de quantização do README
• Há a opinião de que __shfl_down hoje em dia não é recomendado por causa de problemas de sincronização de warp