DeepGEMM: kernels GEMM FP8 limpos e eficientes por meio de scaling granular
(github.com/deepseek-ai)- DeepGEMM é uma biblioteca de kernels tensor core de alto desempenho que reúne em uma única base de código CUDA primitivas computacionais centrais dos LLMs modernos, como GEMM, MoE fused, MQA scoring e HyperConnection
- Todos os kernels são compilados em runtime por um módulo JIT leve; não requer compilação CUDA durante a instalação e exige C++20, CUDA Toolkit, PyTorch e CUTLASS 4.0 ou superior
- Embora aproveite alguns conceitos do CUTLASS e do CuTe, não depende pesadamente de templates nem de estruturas algébricas, e foi projetada para tornar o aprendizado de otimização de kernels NVIDIA GPU mais acessível com um número limitado de funções de kernel centrais
- O escopo de suporte inclui FP8, FP4, BF16 GEMM, grouped GEMM, kernel de MQA logits para DeepSeek v3.2 e Mega MoE, que sobrepõe comunicação e computação; há diferenças nas restrições de layout de memória entre SM90 e SM100
- Mesmo com um design leve, mira desempenho igual ou superior ao de bibliotecas ajustadas por especialistas em diversos formatos de matriz, e inclui uma atualização que atingiu até 1550 TFLOPS no H800
Objetivo e design do DeepGEMM
- DeepGEMM é uma biblioteca de kernels tensor core que integra em uma única base de código CUDA as principais primitivas computacionais usadas em modelos de linguagem de grande porte modernos
- GEMM: FP8, FP4, BF16
- MoE fused com comunicação sobreposta: Mega MoE
- MQA scoring para lightning indexer
- HyperConnection(HC)
- Todos os kernels são compilados em runtime como um módulo Just-In-Time(JIT) leve
- Não é necessária compilação CUDA durante o processo de instalação
- Aproveita alguns conceitos do CUTLASS e do CuTe
- Porém, não depende muito de templates pesados nem de estruturas algébricas
- Mantém a base de código simples ao limitar o número de funções de kernel centrais
- Afirma entregar desempenho equivalente ou superior ao de bibliotecas ajustadas por especialistas em vários formatos de matriz, apesar do design leve
Principais atualizações
- A atualização de 16 de abril de 2026 inclui Mega MoE, GEMM FP8xFP4, FP4 Indexer, PDL, compilação JIT mais rápida e outros itens
- A atualização de 28 de setembro de 2025 adiciona um kernel de scoring weighted ReLU MQA logits para o lightning indexer do DeepSeek v3.2
- Detalhes em #200
- A atualização de 20 de julho de 2025 passa a oferecer suporte tanto a SM90 quanto a SM100, com refatoração completa para um módulo JIT CPP de baixo overhead de CPU
- NVRTC e otimizações SASS pós-compilação foram desativados
- NVRTC aparece como suporte previsto posteriormente
- Como o NVCC 12.9 executa automaticamente o interleaving de FFMA, a otimização pós-compilação não é mais suportada
- Detalhes em #112
- A atualização de 14 de maio de 2025 adiciona kernels de weight gradient para backward denso e MoE
- Detalhes em #95
- A atualização de 7 de maio de 2025 oferece até 10x mais velocidade de compilação com suporte a NVRTC
- Pode ser ativado com
DG_JIT_USE_NVRTC=1 - Em alguns casos, pode haver perda de desempenho
- Detalhes em #94
- Pode ser ativado com
- A atualização de 18 de abril de 2025 atingiu até 1550 TFLOPS no H800
Requisitos e fluxo de instalação
- O ambiente de execução exige uma GPU com arquitetura NVIDIA SM90 ou SM100
- Os requisitos de software são os seguintes
- Python 3.8 ou superior
- Compilador com suporte a C++20
- CUDA Toolkit
- SM90: CUDA 12.3 ou superior
- CUDA 12.9 ou superior é fortemente recomendado para o melhor desempenho
- SM100: CUDA 12.9 ou superior
- PyTorch 2.1 ou superior
- CUTLASS 4.0 ou superior
- Biblioteca
{fmt}
- No ambiente de desenvolvimento, depois de clonar o repositório incluindo submodules,
develop.shfaz os vínculos dos includes necessários e compila o módulo CPP JIT - A instalação consiste em executar
install.she depois importardeep_gemmno projeto Python
Interface GEMM e restrições de layout
- A convenção de nomenclatura dos kernels GEMM do DeepGEMM é
D = C + A @ B - O layout dos shapes de entrada tem NT como referência
fp8_gemm_ntexecutaD = C + A @ B.T
- A implementação SM90 suporta apenas layout de memória NT
- Corresponde à combinação row-major, col-major
- A implementação SM100 suporta todos os layouts de memória NT, TN, NN, TT
- Em ambas as arquiteturas, o scaling factor do LHS deve estar alinhado para TMA e em layout transposto
- SM90 exige o scaling factor no formato FP32
- SM100 exige o formato packed UE8M0, empacotando 4 UE8M0 em um
torch.int
- Operações como transposição de entrada ou casting para FP8 devem ser tratadas separadamente pelo usuário
- A biblioteca oferece funções utilitárias simples em PyTorch, mas elas podem ser lentas
- O foco principal é a otimização de kernels GEMM
GEMM denso e grouped
- O GEMM FP8 non-grouped básico usa as funções
fp8_gemm_{nt, nn, tn, tt} - O grouped GEMM em contiguous layout, diferentemente do grouped GEMM tradicional do CUTLASS, agrupa apenas o eixo M
- N e K devem ser fixos
- É um design voltado a situações em que experts em modelos MoE compartilham o mesmo shape
- No forward pass de treinamento ou no prefilling de inferência, a quantidade de tokens processada por cada expert pode variar
- A forma em que esses tokens são concatenados em um único tensor é chamada de contiguous layout
- Cada segment de expert deve estar alinhado ao GEMM M block size
- O critério de alinhamento é verificado com
get_mk_alignment_for_contiguous_layout()
- Também há uma API grouped no eixo K para backward de weights em MoE
- M e N devem ser fixos
- A função relacionada é
k_grouped_fp8_gemm_tn_contiguous
- Na etapa de decoding de inferência, quando CUDA graph está ativado e a CPU não consegue saber a quantidade de tokens por expert, há suporte a masked grouped GEMM
- Ao fornecer um mask tensor, o kernel calcula apenas a região válida
- A função é
m_grouped_fp8_gemm_nt_masked - Há um exemplo que usa como entrada a saída dos kernels de baixa latência do DeepEP
Kernel MQA para o Indexer do DeepSeek v3.2
- A família de kernels MQA V3.2 oferece uma versão non-paged e uma versão paged
- non-paged é para prefilling
- paged é para decoding
fp8_mqa_logitsrecebe 6 entradasq: tensor E4M3, shape[seq_len, num_heads, head_dim]kv: tensor E4M3 e scaling factor float- O shape do tensor é
[seq_len_kv, head_dim] - O shape do scaling factor é
[seq_len_kv]
- O shape do tensor é
weights: tensor float, shape[seq_len, num_heads]cu_seq_len_k_start,cu_seq_len_k_end: tensor int, shape[seq_len]clean_logits: se deve limpar logits não preenchidos com-inf
- O shape do tensor de saída é
[seq_len, seq_len_kv]e representa logits token-to-token - Cada token
qipercorre os tokenskvjdecu_seq_len_k_start[i]até antes decu_seq_len_k_end[i]- Multiplica
kv_jpelo scaling factor - Calcula valores por head com
q[i, :, :] @ kv_j - Depois aplica ReLU, multiplica por
weights[i, :]e soma para gerar um logit escalar
- Multiplica
- A função da versão paged é
fp8_paged_mqa_logits
Mega MoE
- Mega MoE faz fuse de várias etapas de MoE em um único mega-kernel
- EP dispatch
- linear 1, FP8xFP4
- SwiGLU
- linear 2, FP8xFP4
- EP combine
- Mega MoE sobrepõe comunicação NVLink e computação tensor core
- A execução requer multi-process launch usando symmetric memory
- O fluxo de uso é o seguinte
- Aloca um symmetric memory buffer com
deep_gemm.get_symm_buffer_for_mega_moe- Requer PyTorch 2.9 ou superior
- Transforma os weights, incluindo FP4 e UE8M0 SF, para o layout exigido com
deep_gemm.transform_weights_for_mega_moe - Antes da chamada, copia input, scaling factor, top-k index e top-k weight para o buffer
- Executa o kernel fused mega MoE com
deep_gemm.fp8_fp4_mega_moe
- Aloca um symmetric memory buffer com
- O setup multi-process completo e exemplos de benchmark estão em
tests/test_mega_moe.py
Utilitários e variáveis de ambiente
- As principais funções utilitárias controlam recursos de execução, alinhamento, compilação JIT e conversão de scaling factor
deep_gemm.set_num_sms/get_num_sms: define e consulta o número máximo de SMs a usardeep_gemm.set_tc_util/get_tc_util: define e consulta a taxa aproximada de utilização dos tensor coresdeep_gemm.set_pdl/get_pdl: ativa e desativa Programmatic Dependent Launch(PDL)deep_gemm.set_mk_alignment_for_contiguous_layout/get_mk_alignment_for_contiguous_layout: define e consulta o alinhamento M/K em nível de grupo no contiguous layoutdeep_gemm.transform_sf_into_required_layout: transforma scaling factor para o layout exigidodeep_gemm.get_tma_aligned_size: consulta o tamanho de alinhamento TMA necessário
- Variáveis de ambiente relacionadas ao JIT controlam saída de debug, local do cache, escolha do compilador e opções de profiling
DG_JIT_DEBUG: imprime informações de debug do JITDG_PRINT_CONFIGS: imprime a config escolhida por shapeDG_JIT_CACHE_DIR: diretório de cache dos kernels compilados; o padrão é$HOME/.deep_gemmDG_JIT_USE_NVRTC: usa NVRTC em vez de NVCC, permitindo compilação rápida, mas em alguns casos o desempenho pode ser menorDG_JIT_NVCC_COMPILER: caminho do compilador NVCCDG_JIT_CPP_STANDARD: versão do padrão C++; o padrão é20
- Também são fornecidas variáveis de ambiente para debug e profiling
DG_JIT_DUMP_ASM,DG_JIT_DUMP_PTX,DG_JIT_DUMP_SASS: despejam a saída PTX e SASSDG_JIT_WITH_LINEINFO: inclui informações de linha de origem para ferramentas de profilingDG_COMM_KERNEL_DEBUG: inicializa o symmetric buffer com 0 antes da chamada ao Mega MoEDG_USE_NVIDIA_TOOLS: ignora o profiling interno ao executar ferramentas NVIDIA externas
- Opções de build controlam instalação e forma de carregamento dos kernels
DG_SKIP_CUDA_BUILD: pula o build da extensão CUDA durante a instalaçãoDG_FORCE_BUILD: força build local em vez de baixar uma wheel pre-builtDG_JIT_USE_RUNTIME_API: usa CUDA Runtime API para carregar kernels; requer CUDA runtime 12.8 ou superior
Licença e citação
- O repositório DeepGEMM é disponibilizado sob a MIT License
- O projeto afirma ter sido inspirado no CUTLASS
- O título do item de citação é
DeepGEMM: clean and efficient BLAS kernel library on GPU
1 comentários
Opiniões no Hacker News
O interleaving de FFMA em SASS parece realmente impressionante
Ao perceberem que o desempenho do kernel FP8 do CUTLASS melhorou entre o NVCC 12.2 e 12.3 e compararem o SASS compilado, aparentemente descobriram que um bit em várias instruções FADD havia sido invertido em um padrão intercalado e, consultando uma implementação open source de um assembler CUDA, identificaram que esse bit era o bit de yield, que faz a warp atual ceder para que outra warp seja executada
É impressionante que tenham criado um script para modificar as instruções FFMA no binário compilado usando isso e, como quando uma warp cede não é possível reutilizar registradores, também inverteram o bit de reuse, conseguindo sobrepor melhor as instruções MMA e as instruções FFMA de promoção em GEMM FP8 com escalonamento fino, elevando o desempenho em mais de 10% em alguns casos
Só que, nesse problema específico, parece que outras empresas de IA ainda não tinham sentido a necessidade de aplicá-la, e é provável que todos acabem chegando a pontos parecidos
Casos assim mostram bem o quanto os compiladores atuais ainda estão longe de extrair desempenho do hardware apenas a partir de código de alto nível
Fico curioso sobre o que seria necessário para que técnicas tradicionais de compiladores ou agentes de otimização baseados em IA chegassem a esse tipo de resultado
Os números de aceleração relatados são comparados com a própria linha de base baseada em CUTLASS
Fico curioso se alguém fez uma comparação direta de desempenho com o cuBLAS
Os resultados de GEMM com CUTLASS que vi até agora ficavam, em geral, dentro de algo como 10% em relação ao cuBLAS; se a melhoria de 2x a 2,5x mencionada no artigo se mantiver, seria realmente impressionante
Para começar, o cuBLAS precisa da API estendida cuBLASLt para lidar com operações de precisão mista como FP8
Além disso, ele não oferece suporte a combinações de tipos que pareceriam adequadas, como E5M2 x E5M2 em A x B, mas oferece suporte a E5M2 x E4M3; e as restrições continuam, como a exigência de que, em Ampere, Hopper e Blackwell, a matriz A esteja sempre em layout transposto
Integrei um benchmark de FP8 com cuBLASLt ao meu repositório "Less Slow C++" <https://github.com/ashvardanian/less_slow.cpp> e o adicionei à lista de benchmarks existentes de cuBLAS e de CUDA/PTX escritos por mim
Estou executando em uma GPU H200, que deve ter o mesmo desempenho de uma H100, e, com entradas quadradas, a vazão atinge um pico de cerca de 1,35 Peta-ops
256 ficou em 2,68T/s; 512, em 20,49T/s; 1024, em 144,23T/s; 2048, em 665,68T/s; 4096, em 1,26P/s; 8192, em 1,34P/s; e 16384, em 1,23P/s, o que corresponde a cerca de 67% do número que a NVIDIA divulga para GEMM denso <https://resources.nvidia.com/en-us-data-center-overview-mc/e...>
Eu havia imaginado que a linha de base escolheria o melhor entre cuBLAS e CUTLASS
Esse tipo de open source mostra muito bem o objetivo de alcançar eficiência na indústria
Ainda assim, o benefício desse software provavelmente irá mais para grandes empresas que servem modelos em larga escala, ou seja, potenciais concorrentes da DeepSeek, do que para a comunidade open source em geral, que quer aprender, experimentar ou servir modelos em hardware de consumidor
Não sei bem se otimizar rumo a precisões cada vez mais baixas é bom no longo prazo
Isso significa que os modelos são, na prática, bastante esparsos; e, embora isso possa ser verdade hoje, acho mais provável que seja porque há ideias ruins misturadas na forma de treinamento, não porque eles precisem ser inerentemente tão esparsos
Tornar possível treinar modelos realmente bons apenas em precisão mais alta é um problema de pesquisa; treinamento e inferência em baixa precisão são problemas de engenharia
Fazemos esse tipo de coisa desde a época das CNNs, pelo menos há 9 anos, e acho que ainda teremos mais alguns anos disso pela frente
Isso pode acabar se tornando irrelevante por causa do MXFP, o suporte nativo a microscaling do Blackwell
No Hopper, isso foi basicamente implementado manualmente em uma granularidade mais grossa, mas usando coeficientes de escala FP32
Demonstrações públicas de alta qualidade como essa mostram bem onde está o moat da $NVDA
GPUs de uso geral são muito flexíveis, então é possível programá-las para realizar várias tarefas que fazem sentido, mas que o fornecedor do hardware talvez não tenha imaginado desde o início
Porém, se a previsão for de que o futuro vai convergir cada vez mais para suporte de hardware dedicado, eliminando espaço para esse tipo de otimização de software, o chamado moat do CUDA desmorona
Para continuar nesse jogo, a NVIDIA está, de certa forma, derrubando o próprio moat :p
Uau, é licença MIT
Seria bom se as grandes empresas adotassem esse modelo de colaboração open source
Continuo me perguntando por que existem instruções não documentadas
Acho que, mesmo que não sejam perfeitamente estáveis, seria melhor disponibilizá-las aos usuários
Esse tipo de informação provavelmente é documentado internamente, então não entendo por que não a tornam pública
Segurança baseada em obscuridade não funciona, e os concorrentes vão fazer engenharia reversa de tudo de qualquer jeito
Pode ser por falta de tempo, ou porque não querem sugerir suporte a recursos instáveis ou experimentais
Se o impacto ficar restrito mais ou menos à equipe ao lado, também fica muito mais fácil mudar
Provavelmente só está em algum documento de projeto da arquitetura ou especificação, e é claro que eles não gostariam de compartilhar esse tipo de documento
Sinceramente, é um conteúdo que vai além do meu escopo de uso e compreensão
Ainda assim, é realmente gratificante e revigorante ver esse tipo de descoberta e melhoria sendo compartilhado para que todos possam se beneficiar
Ela é muito importante em multiplicação de matrizes e cargas de trabalho de deep learning
No SASS da NVIDIA, a instrução FFMA é codificada como uma instrução de 64 ou 128 bits e tem vários bits de controle que determinam seu comportamento exato
Quando o bit de yield é definido, ele informa ao escalonador de warps que, após essa instrução, o warp atual pode ceder a execução, e o hardware pode executar outro warp para ocultar a latência
A GPU obtém alta vazão por meio de paralelismo massivo e, se um warp para por espera de memória ou algo parecido, outro warp pode avançar
O bit de reuse indica se o registrador de origem pode ser reutilizado logo na operação seguinte e, se o bit de yield estiver definido, ele precisa obrigatoriamente ser desligado
Isso porque, quando um warp cede a execução, o próximo a executar pode não ser esse warp, e outro warp pode alterar o estado do arquivo de registradores, de modo que o hardware não consegue garantir que os valores dos registradores sejam preservados para além do yield
Ao definir os bits de yield nas instruções FFMA em um padrão alternado, o compilador cria pontos explícitos de escalonamento onde outros warps podem avançar e, para manter a correção, também precisa limpar os bits de reuse dessas instruções
Essa mudança ajuda especialmente a sobrepor as instruções MMA, que são o núcleo da multiplicação de matrizes, com as instruções FFMA de promoção, que fazem a conversão para acumular FP8 com maior precisão
FP8 GEMM normalmente precisa converter para uma precisão maior para a acumulação e depois converter de volta, o que gera FFMAs adicionais; isso reduz a demanda por largura de banda de memória, mas cria um padrão de computação complexo, misturado com operações de promoção/rebaixamento
“Escalonamento fino” parece se referir ao trabalho de gerenciar cuidadosamente a precisão em vários pontos do cálculo
A manipulação do bit de yield faz com que as operações de cálculo e as conversões de formato fiquem melhor intercaladas, permitindo usar as unidades de execução da GPU com mais eficiência; sem essa otimização, o escalonador de warps pode não encontrar oportunidades naturais de troca, e os recursos de computação podem acabar sendo subutilizados