Compilando LLMs em um MegaKernel para viabilizar inferência de baixa latência
(zhihaojia.medium.com)- Foi desenvolvido um compilador que converte automaticamente a inferência de LLM em um único megakernel
- A abordagem MegaKernel (kernel persistente) integra completamente computação e comunicação da inferência de LLM em um único kernel de GPU, possibilitando latência muito baixa
- Existe o problema de que, devido à estrutura distribuída dos frameworks de ML e das bibliotecas de kernel existentes, é muito difícil transformar todo o pipeline em um único kernel
- O Mirage Persistent Kernel (MPK) converte automaticamente a inferência de LLM em múltiplas GPUs em um megakernel de alto desempenho por meio de um compilador e sistema de runtime
- O MPK transforma o grafo computacional em um grafo de tarefas granular, maximizando software pipelining e a sobreposição entre computação e comunicação
- Com o MPK, a latência de geração de tokens diminui em comparação com sistemas existentes, e o ganho de desempenho cresce ainda mais à medida que o número de GPUs aumenta
Visão geral e vantagens da abordagem MegaKernel
- Na inferência de grandes modelos de linguagem (LLMs), uma das formas mais eficazes de reduzir a latência é fundir todos os processos de computação e comunicação em um único megakernel (kernel unificado)
- Nessa abordagem, um único kernel de GPU executa sem interrupção todo o processamento, incluindo operações por camada de todo o modelo e a comunicação entre GPUs
- Os principais benefícios são os seguintes
- Eliminação do overhead de lançamento de kernel ao dispensar chamadas repetidas de kernel
- Possibilidade de realizar software pipelining ao longo das camadas
- Ocultação de latência ao executar computação e comunicação simultaneamente
Limitações anteriores e o surgimento do MPK
- Frameworks de ML como PyTorch, Triton e TVM não oferecem, por natureza, suporte à geração automática end-to-end de megakernel
- Sistemas reais de LLM são compostos por combinações de várias bibliotecas de kernel, como NCCL/NVSHMEM (comunicação), FlashInfer/FlashAttention (atenção), CUDA/Triton (operações customizadas), o que dificulta a integração em um único kernel
- Nesse contexto, pesquisadores da CMU, UW, Berkeley, NVIDIA e Tsinghua desenvolveram o Mirage Persistent Kernel (MPK)
- O MPK combina compilador e runtime para converter automaticamente todo o pipeline de inferência de LLM em um megakernel de alto desempenho
Valor central do MPK
- O MPK elimina completamente o overhead de lançamento de kernel e maximiza a sobreposição entre computação, carregamento de dados e comunicação entre camadas, viabilizando um ambiente de inferência de LLM de latência ultrabaixa
- Em testes reais (prompt de 39 tokens, geração de 512 tokens, sem speculative decoding),
- Em um ambiente com uma única GPU NVIDIA A100 40GB, em comparação com a latência de decodificação por token (14,5 ms) de sistemas otimizados existentes como vLLM/SGLang, o MPK reduziu para 12,5 ms
- Esse número se aproxima do limite inferior teórico (10 ms), considerando largura de banda de memória de 1,6 TB/s e carregamento de 16 GB de pesos
- Ao integrar completamente computação e comunicação em ambientes com múltiplas GPUs, a vantagem de desempenho do MPK se destaca ainda mais conforme aumenta o número de GPUs
Estrutura de funcionamento do MPK em detalhe
Parte 1. Compilador – conversão do grafo computacional de LLM em grafo de tarefas
- Em geral, as operações de LLM são representadas como um grafo computacional em que cada operação (por exemplo, multiplicação de matrizes, atenção) ou operação de comunicação (por exemplo, all-reduce) é um nó, e as dependências de dados são as arestas
- No desenho tradicional, é comum executar um kernel separado por operador, mas isso reflete apenas dependências no nível de kernel, e não no nível real dos dados dependentes, o que limita as oportunidades de pipelining
- Exemplo: quando há um all-reduce após uma multiplicação de matrizes, o all-reduce só começa depois que toda a multiplicação termina. Na prática, é possível dividir os dados e explorar execução parcial e relações de dependência
- O compilador do MPK refina o grafo computacional e o converte automaticamente em um fine-grained task graph adequado às unidades reais de dados
- Cada tarefa (retângulo) é uma unidade de computação/comunicação atribuída a um SM individual da GPU
- Cada evento (círculo) é um ponto de sincronização entre tarefas
- As arestas entre tarefas e eventos expressam com eficiência as dependências de dados e de controle
- Graças a esse grafo de tarefas, o MPK permite sobrepor computação e comunicação de forma mais parcial ou paralela
- O Mirage kernel superoptimizer também gera automaticamente implementações CUDA de alto desempenho adequadas para cada tarefa
Parte 2. Runtime – execução do grafo de tarefas dentro do megakernel
- O runtime do MPK executa completamente o grafo de tarefas apenas dentro de um único kernel da GPU (megakernel)
- Todos os SMs (Streaming Multiprocessors) da GPU são divididos estaticamente entre os papéis de worker e scheduler
Worker
- Cada worker opera no nível de SM e gerencia sua própria fila de tarefas dedicada
- Em loop, ele
- Busca a próxima tarefa na fila
- Executa a tarefa (por exemplo, matmul, attention, transferência de dados)
- Ao concluir, notifica um evento
- Repete o processo
- Isso permite otimizar o uso de recursos de cada worker e realizar operações assíncronas entre camadas
Scheduler
- Um scheduler distribuído opera no nível de um único warp em cada SM, com até 4 schedulers podendo rodar simultaneamente
- Cada scheduler gerencia uma fila de eventos ativados e atribui aos workers as tarefas cujas condições foram satisfeitas
- Assim, torna-se possível distribuir um grande volume de tarefas sem overhead de sincronização centralizada
Método de execução baseado em eventos
- Quando uma tarefa é concluída, ela incrementa um contador de evento específico. Quando o contador atinge um limiar, o evento é ativado e inserido na fila do scheduler
- O scheduler então executa as tarefas subsequentes que têm relações de dependência com esse evento
- Com isso, fine-grained software pipelining e sobreposição entre computação e comunicação acontecem de forma natural
- Exemplo: o matmul de uma camada e a attention de outra camada podem ser executados ao mesmo tempo
- Assim que resultados parcialmente concluídos de um matmul aparecem, já é possível iniciar a comunicação all-reduce
- Como todo o agendamento e a troca entre tarefas ocorrem dentro de um único contexto de kernel, o overhead entre tarefas fica em torno de 1–2 microssegundos (μs), extremamente baixo
Direções futuras
-
Objetivo do MPK: permitir que desenvolvedores, escrevendo apenas uma pequena quantidade de código Python (dezenas de linhas), consigam compilar facilmente LLMs em megakernels e extrair o máximo desempenho
-
Principais direções de evolução
- Suporte a arquiteturas de GPU mais recentes: por exemplo, foco na NVIDIA Blackwell e em abordagens especializadas no nível de warp
- Tratamento de workloads dinâmicos: pesquisa de estratégias de compilação para modelos que exigem fluxo de controle dinâmico, como mixture-of-experts (MoE)
- Agendamento avançado de tarefas: pesquisa sobre políticas modernas, como priorização e otimização de throughput, e busca por possibilidades de aplicação
-
O MPK apresenta um ponto de inflexão fundamental na forma de compilar e executar workloads de inferência de LLM em GPU, e busca ampliar a colaboração com a comunidade
Materiais adicionais
- O código e a documentação do MPK (Mirage Persistent Kernel), bem como os resultados mais recentes da pesquisa, podem ser consultados no GitHub (https://github.com/mirage-project/mirage)
1 comentários
Comentários no Hacker News
Para o autor: é interessante ver que a abordagem de intérprete on-GPU parece uma direção futura muito promissora. Há também outro trabalho com uma abordagem quase idêntica, então recomendo ver este post relacionado. O modelo fundamental de programação do CUDA (por exemplo, lançamento de kernel) está sendo contornado para paralelização baseada em tarefas bem granulares, e pude ver em primeira mão que isso aumenta o aproveitamento do hardware. Fico me perguntando se o CUDA não estava nos limitando de várias formas. Também tenho expectativa sobre a possibilidade de essa pesquisa do autor entrar como backend experimental do PyTorch. E, por fim, uma pequena observação de typo: os dois parágrafos da primeira parte estão quase idênticos.
Trabalhei de perto por um bom tempo com vLLM e SGLang, e estou convencido de que este projeto é exatamente a forma ideal do próximo projeto. Foi impressionante ver a análise do grafo de dependência das operações e como vocês fundem operações ou fazem um agendamento de tarefas mais inteligente. Parabéns à equipe.
Dei uma olhada no post e no README do GitHub e achei o projeto realmente muito legal. Fiquei curioso se esse tipo de otimização poderia ser aplicado não só à inferência, mas também à etapa de treinamento. Em especial, entendo que a fusão de operações de backward e da comunicação de gradientes é um desafio. Pelo que sei, no momento vocês ainda não suportam workloads dinâmicos (por exemplo, MoE), e recentemente saiu um paper sobre processar MoE em um único kernel: FlashDMoE: Fast Distributed MoE in a Single Kernel.
Obrigado por ler o post e até o README. Também é possível dar suporte à etapa de treinamento, mas em geral os kernels de treinamento são maiores, então o overhead de lançamento de kernel não costuma ser um problema tão grande; por isso, a inferência (especialmente a de baixa latência) é a principal beneficiada. Também achei interessante o paper FlashDMoE que você compartilhou, e quero destacar que dar suporte a modelos MoE é um dos próximos objetivos.
Pessoalmente, tenho uma visão um tanto cética sobre investir tempo em otimização de treinamento baseado em gradiente. Na prática, muitas tarefas de treinamento têm características de valores discretos, e acho que o aprendizado baseado em gradiente não lida bem com isso.
O próximo passo é claramente compilar direto para Verilog e comprar hardware de LLM no AliExpress.
Compartilhando um texto que apresenta tecnologias de hardware como Chisel. Antes do surgimento de IA e GPUs, essa ideia de converter software diretamente em hardware era uma abordagem promissora. Como o avanço das CPUs está estagnado, o desejo de otimizar mais a camada intermediária entre software e hardware continua forte, mas a computação paralela no estilo GPU provavelmente seguirá como a principal forma de aceleração. CPUs de uso geral devem acabar ficando no papel de pequenos cérebros que gerenciam GPUs. Ainda assim, minha expectativa é que uma abordagem de converter software diretamente em hardware dificilmente se torne dominante.
Se a estrutura dos LLMs se estabilizar daqui a 5~10 anos, talvez passe a fazer sentido mapeá-los diretamente em hardware. Com a tecnologia atual, pode ser possível colocar até modelos com dezenas de bilhões de parâmetros em um único wafer usando apenas portas lógicas de ultra baixa precisão, próximas de 1,5 bit. À medida que a precisão aumenta, o número de portas cresce exponencialmente, então, por enquanto, é mais eficiente manter a memória dos pesos e compartilhar unidades de computação. No futuro, desenvolver LLMs de ultra baixa precisão será uma tarefa essencial.
Brincadeira sobre como o custo de treinamento já é alto e, se ainda somar o custo da máscara, a situação só piora; e a avaliação mais fria de que, na prática, startups de hardware para IA já vêm tentando algo nessa direção há muito tempo.
Se algo como um LLM-in-a-box realmente existir, isso seria bastante atraente. Em breve talvez eu tenha a chance de trabalhar em um ambiente offline (
air-gap), e uma solução dessas parece que seria muito útil.Rodei o código diretamente em um ambiente de GPU da Modal, e os ganhos de desempenho alegados na pesquisa realmente se reproduzem na prática. Compartilho o código com os resultados do projeto mirage. Na combinação Triton + FlashInfer, a latência ficou em cerca de 19,2 ms por token; com MPK, nas mesmas condições, tive uma grande melhora para 7,7 ms.
Já participei de uma pequena competição de CUDA. Era sobre algoritmos paralelos para imagem ou visão, e, para parecer esperto, eu fazia cache dos resultados intermediários na memória. Quando vi o resultado da competição, fiquei surpreso que outras pessoas tinham enviado código muito mais rápido que o meu. A razão era que, em vez de fazer cache desses resultados intermediários, eles simplesmente recalculavam tudo o tempo todo. O custo computacional era muito menor que o custo de ida e volta à memória. Suspeito que este projeto seja algo parecido. Ao compilar em megakernel, as fronteiras entre camadas desaparecem, então o compartilhamento de resultados intermediários diminui e a quantidade de computação aumenta, mas, no geral, a grande vantagem vem da redução do tráfego de memória. Especialmente em redes convolucionais, deve haver algum sweet spot, mas não sei como o megakernel lida com isso.
Continuam surgindo novas metáforas para LLMs. Fiquei pensando se talvez possamos tratar LLMs como transistores. Imagino que agora estamos numa fase parecida com a de computadores do tamanho de uma sala que só faziam multiplicação com cartões perfurados. É divertido imaginar o que aconteceria se fosse possível rodar 1 milhão de consultas do o3-pro ao mesmo tempo.
Este projeto vem da CMU (Carnegie Mellon). Também houve menção ao blog da Hazy Research, de Stanford, sobre megakernels, No Bubbles. É impressionante ver quão ativa está a competição nessa área. (Adendo) Também existe um paper sobre o quadro maior do projeto "mirage", mas ele não aborda a abordagem de megakernel: link do paper
O próprio autor do post respondeu diretamente. Concordo que a pesquisa com Stanford está acontecendo em paralelo. A principal diferença é que estamos focados em um compilador para geração automatizada de megakernels.
Também vale mencionar que o ThunderKittens, da Hazy Research, é uma biblioteca muito legal. Recentemente, tem havido um grande esforço concentrado em formalização, pipelining, divisão e conquista, maximização de eficiência e desenvolvimento de compiladores/DSLs dedicados para aproveitar ao máximo os modelos recentes de GPU da NVIDIA.
Os números de desempenho do Qwen 8B, se forem confirmados, são bem impressionantes. Parece mais prático do que abordagens anteriores com megakernel. Esse tipo de kernel mantido um por SM me lembra o antigo Larrabee. Fico curioso sobre como o mundo seria hoje se, em vez do CUDA atual, tivéssemos seguido o caminho mais tradicional de processo-thread-SIMD.
Uma ideia sobre criar LLMs fixos em ASIC puro, em vez de inferência baseada em software. Haveria vantagem de custo? Seria possível fornecer camadas que ainda pudessem ser tratadas ou ajustadas por software? Como já estamos quase chegando a um nível "bom o suficiente", talvez em 2~4 anos as pessoas decidam usar chips especializados congelados por um bom tempo. Fico me perguntando em que momento exatamente os benefícios desse hardware ultraespecializado começariam a se destacar.