DeepGEMM: kernel GEMM FP8 limpo e eficiente com escalonamento fino

DeepGEMM

DeepGEMM é uma biblioteca para multiplicação de matrizes gerais FP8 (GEMM) que oferece suporte ao escalonamento fino proposto no DeepSeek-V3. A biblioteca oferece suporte a GEMM agrupado comum e para Mix-of-Experts (MoE), é escrita em CUDA e não exige compilação durante a instalação. Ela oferece suporte aos tensor cores NVIDIA Hopper e usa acumulação em duas etapas com CUDA cores para resolver a imprecisão da acumulação FP8 em tensor cores. Aproveita parcialmente conceitos do CUTLASS e do CuTe, mas mantém a simplicidade ao minimizar a dependência de templates ou álgebra. Com uma única função de kernel central de cerca de 300 linhas de código, é um bom recurso para aprender sobre multiplicação de matrizes FP8 no Hopper e técnicas de otimização. Apesar do design leve, iguala ou supera o desempenho de bibliotecas ajustadas por especialistas em diversos formatos de matriz.

Desempenho

Todos os formatos que podem ser usados na inferência do DeepSeek-V3/R1 foram testados no H800 SXM5 com NVCC 12.8. Todas as métricas de ganho de velocidade foram calculadas em comparação com uma implementação otimizada internamente baseada no CUTLASS 3.6. Em alguns formatos, o desempenho pode não ser bom, e PRs de otimização são bem-vindos.

GEMM comum (modelo denso)

• A medição do desempenho do DeepGEMM em vários tamanhos de matriz mostrou ganho de velocidade de até 2,7x em determinados tamanhos.

GEMM agrupado para modelos MoE (layout contínuo)

• Dependendo do número de grupos e do tamanho das matrizes de cada grupo, mostrou ganho de velocidade de até 1,2x.

GEMM agrupado para modelos MoE (layout com máscara)

• Com uso de layout com máscara, mostrou ganho de velocidade de até 1,2x.

Início rápido

Requisitos

• GPU com arquitetura Hopper, com suporte a sm_90a
• Python 3.8 ou superior
• CUDA 12.3 ou superior (12.8 ou superior recomendado para melhor desempenho)
• PyTorch 2.1 ou superior
• CUTLASS 3.6 ou superior

Desenvolvimento

• Explicação do processo de desenvolvimento, incluindo clonagem de submódulos, criação de link simbólico, compilação JIT e teste de todas as implementações de GEMM.

Instalação

• deep_gemm pode ser importado e usado em projetos Python.

Interface

Observações

• Esta biblioteca inclui apenas kernels GEMM e oferece suporte apenas ao formato NT. Operações de transposição ou outras tarefas de casting FP8 devem ser implementadas de forma independente.

GEMM denso comum (não agrupado)

• Fornece funções para executar GEMM FP8 básico não agrupado.

GEMM agrupado (layout contínuo)

• Projetado para cenários em modelos MoE nos quais os especialistas compartilham o mesmo formato.

GEMM agrupado (layout com máscara)

• Na etapa de decodificação da inferência, fornece um tensor de máscara para calcular apenas as partes válidas.

Utilitários

• Fornece várias funções utilitárias e variáveis de ambiente para ajudar na otimização de desempenho.

Otimização

Especialização persistente de warp

• Segue o design do CUTLASS, sobrepondo movimentação de dados, instruções tensor core MMA e promoção com CUDA cores.

Recurso TMA do Hopper

• Usa TMA para acelerar a movimentação de dados.

Otimizações detalhadas comuns

• Melhora o desempenho por meio de várias técnicas de otimização.

Agendador de blocos unificado e otimizado

• Fornece um agendador para todos os kernels não agrupados e agrupados.

Design totalmente JIT

• Melhora o desempenho com um design JIT que não exige compilação durante a instalação.

Tamanhos de bloco não alinhados

• Oferece suporte a tamanhos de bloco não alinhados para maximizar a utilização de SM em determinados formatos.

Intercalação FFMA SASS

• Modifica instruções FFMA para melhorar o paralelismo em nível de warp e aumentar o desempenho.

Agradecimentos

• O DeepGEMM foi inspirado no projeto CUTLASS, e expressa agradecimento e respeito aos desenvolvedores.

Licença

• Disponibilizado sob a licença MIT.