Resumo Técnico: MPK - Um Compilador e Runtime para Programas de Tensores Mega-Kernelizados

1. Declaração do Problema

Os atuais sistemas de inferência de GPU de alto desempenho para modelos de linguagem de grande escala (LLMs) e outros programas de tensores dependem predominantemente de um modelo de execução de kernel por operador. Neste paradigma, cada operador (por exemplo, multiplicação de matrizes, atenção, all-reduce) é executado como um kernel de GPU distinto. Embora esta abordagem simplifique a programação, ela introduz três limitações críticas que prejudicam o desempenho:

1. Barreiras de Kernel Implícitas: A execução convencional de GPU impõe uma barreira entre lançamentos de kernels consecutivos na mesma stream. Isso impede o pipeline de software entre operadores (cross-operator software pipelining), forçando operadores dependentes a executar estritamente em sequência e criando "bolhas de pipeline" onde os recursos de hardware (Tensor Cores, TMAs) permanecem ociosos.
2. Sobreposição de Grão Grosso: As dependências são impostas na granularidade de operadores inteiros. Consequentemente, a sobreposição fina entre computação e comunicação é impossível. Por exemplo, uma operação de AllReduce deve esperar pela conclusão da totalidade de um MatMul precedente, mesmo que fragmentos específicos da redução dependam apenas de fragmentos específicos da saída da multiplicação.
3. Sobrecarga de Lançamento e Flexibilidade Estática: Executar centenas ou milhares de kernels por iteração de inferência acarreta uma sobrecarga de lançamento significativa. Embora os CUDA Graphs mitiguem isso, eles são amplamente estáticos; mudanças dinâmicas no fluxo de controle, formas de tensores ou dependências de dados frequentemente exigem a reinstanciação do grafo, limitando a flexibilidade para cargas de trabalho dinâmicas, como o serviço de LLMs.

Soluções existentes, como mega-kernels manuais (ex: FlashDMoE, Spector et al.), oferecem benefícios de desempenho, mas exigem um esforço de engenharia extensivo e profundo conhecimento de GPU, falhando em generalizar entre modelos ou hardware.

2. Metodologia

Os autores apresentam o Mirage Persistent Kernel (MPK), o primeiro sistema de compilador e runtime que transforma automaticamente a inferência de modelos multi-GPU em um único mega-kernel de alto desempenho. O MPK alcança isso através de dois componentes principais:

A. Representação de Grafo ao Nível de SM (tGraph)

O MPK introduz uma nova representação de grafo chamada tGraph que captura computação e comunicação na granularidade de Streaming Multiprocessors (SMs) individuais, em vez de todo o GPU.

• Nós: Representam tarefas (computação ou comunicação) executadas em um único SM.
• Arestas: Codificam dependências finas entre tarefas via eventos.
• Vantagem: Esta granularidade expõe o paralelismo inacessível aos modelos de nível de kernel, permitindo o software pipelining entre operadores e a sobreposição fina de computação e comunicação.

Otimizações do Compilador:
O compilador MPK transforma um programa de tensor padrão em um tGraph otimizado através de várias etapas:

1. Decomposição de Operador: Divide operadores em conjuntos disjuntos de tarefas baseados na partição do tensor de saída (ex: tiling de linhas/colunas de multiplicação de matrizes) para corresponder ao número de SMs disponíveis.
2. Análise de Dependência: Enumera pares produtor-consumidor entre tarefas para inserir eventos de sincronização apenas onde existem sobreposições de dados.
3. Fusão de Eventos: Mescla pontos de sincronização redundantes usando fusão de conjunto de sucessores (mesclando eventos com conjuntos de consumidores idênticos) e fusão de conjunto de predecessores (mesclando eventos com produtores idênticos) para reduzir a complexidade do grafo.
4. Normalização: Transforma o grafo para que cada tarefa tenha no máximo um evento dependente e um evento de gatilho. Isso é alcançado inserindo tarefas e eventos fictícios para lidar com fan-in/fan-out, garantindo uma representação uniforme e de baixa sobrecarga.
5. Linearização: Utiliza um algoritmo baseado em BFS para atribuir índices contíguos a tarefas disparadas pelo mesmo evento. Isso permite que o runtime codifique o fan-out de eventos como um simples intervalo de índice (primeiro/último ID de tarefa) em vez de uma lista explícita, reduzindo significativamente a pegada de memória do dispositivo.
6. Geração de Implementação de Tarefa: Utiliza superotimização de compilador (Mirage) para gerar automaticamente implementações CUDA de alto desempenho para cada tarefa individual, incluindo otimizações intra-SM como reutilização de registradores e layout de memória compartilhada.

B. Runtime Paralelo In-Kernel

O MPK executa o tGraph dentro de um único mega-kernel persistente usando uma arquitetura descentralizada e orientada a eventos:

• Particionamento Worker/Scheduler: Os SMs do GPU são estaticamente particionados em Workers (que executam tarefas) e Schedulers (que gerenciam filas de eventos).
• Execução Orientada a Eventos: Os Schedulers monitoram a ativação de eventos (onde todos os pré-requisitos foram atendidos) e despacham as tarefas correspondentes para os workers. Isso elimina a necessidade de sincronização do lado do host ou lançamentos repetidos de kernel.
• Lançamento de Tarefas Híbrido: Para equilibrar latência e balanceamento de carga, o MPK emprega dois modos:
  • Just-In-Time (JIT): As tarefas são lançadas apenas após seus critérios serem atendidos. Isso se adapta a cargas de trabalho dinâmicas (ex: comprimentos de sequência variáveis em atenção) para evitar desbalanceamento de carga.
  • Ahead-Of-Time (AOT): As tarefas são pré-enfileiradas nas filas dos workers antes que suas dependências sejam atendidas. Isso reduz a latência de despacho para operadores previsíveis.
• Memória Compartilhada Paginada: Para permitir o software pipelining entre tarefas, o MPK introduz uma abstração paginada para a memória compartilhada. As tarefas adquirem e liberam páginas de tamanho fixo, permitindo que a próxima tarefa comece a fazer o prefetch de dados para a memória compartilhada enquanto a tarefa atual computa, sem conflitos de recursos.

3. Principais Contribuições

1. Primeiro Sistema de Mega-Kernel Automatizado: O MPK é o primeiro sistema que transforma automaticamente programas de tensores multi-GPU em um único mega-kernel, removendo a necessidade de fusão de kernel manual.
2. Abstração ao Nível de SM: A introdução da representação tGraph desloca o rastreamento de dependências do nível de operador para o nível de tarefa-SM, permitindo otimizações (pipelining, sobreposição fina) anteriormente inviáveis sob o modelo de kernel por operador.
3. Compilador e Runtime End-to-End: O sistema integra um compilador que otimiza grafos de tarefas e gera código CUDA com um runtime que executa essas tarefas assincronamente dentro de um único kernel, suportando cargas de trabalho dinâmicas como continuous batching e paged attention.
4. Estratégia de Execução Híbrida: A combinação de lançamento de tarefas JIT e AOT permite que o sistema lide tanto com a natureza dinâmica da inferência de LLM (variáveis de batch/sequência) quanto com a eficiência estática de operadores fixos.

4. Resultados de Avaliação

Os autores avaliaram o MPK em cinco LLMs amplamente utilizados (incluindo variantes de Qwen3 e Llama-3) em três gerações de GPUs NVIDIA (A100, H100, B200), comparando-o com sistemas de estado da arte de kernel por operador (SGLang, vLLM) e o PyTorch vanilla.

• Ganhos de Desempenho: O MPK alcança de 1.0× a 1.7× de menor latência de inferência end-to-end em comparação com SGLang e vLLM. As melhorias são mais pronunciadas para modelos menores e gerações de hardware mais novas (ex: B200), onde a sobrecarga de lançamento de kernel e as bolhas de pipeline se tornam mais significativas em relação ao tempo de computação.
• Throughput: Em configurações multi-GPU (até 8 H100s), o MPK melhora o throughput em 1.1× a 1.4× sobre as bases otimizadas.
• Limites de Latência: Em um A100, o MPK reduziu a latência de decodificação por token do Qwen3-8B de 14.5 ms (vLLM/SGLang) para 12.5 ms, aproximando-se do limite inferior de hardware estimado de ~10 ms.
• Estudos de Ablação:
  • Cross-Task Pipelining: Reduziu o tempo de execução da tarefa em 1.2–1.3×.
  • Sobreposição Computação-Comunicação: Reduziu a latência por iteração em 1.1×.
  • Redução de Lançamento de Kernel: Eliminou a sobrecarga de ~1.1 ms por token associada aos 293 lançamentos de kernel por token nos sistemas de base.
• Escalabilidade: O MPK escala efetivamente através de 2, 4 e 8 GPUs, mantendo acelerações sobre as bases.

5. Significância e Alegações

O artigo afirma que o MPK representa uma mudança de paradigma nos sistemas de inferência de GPU ao unificar o pipeline de execução em um único mega-kernel. Sua significância reside em:

• Quebrar a Barreira do Kernel: Ao mover a sincronização para o nível de SM, o MPK desbloqueia o potencial de utilização de hardware que é atualmente inacessível para sistemas de kernel por operador.
• Automação vs. Esforço Manual: Ao contrário de abordagens anteriores de mega-kernel que exigiam design manual especializado, o MPK automatiza o processo, tornando a execução de mega-kernel de alto desempenho acessível a desenvolvedores padrão de PyTorch com mudanças mínimas de código.
• Proximidade ao Hardware: O sistema empurra o desempenho da inferência de LLM próximo aos limites teóricos do hardware subjacente, minimizando sobrecargas de software (lançamento, agendamento, sincronização) e maximizando a sobreposição de hardware.

Os autores concluem que o MPK abre um novo caminho para sistemas de inferência de alto desempenho impulsionados por compiladores, que preservam a flexibilidade dos modelos de programação existentes enquanto entregam desempenho próximo ao limite do hardware.