Hybrid JIT-CUDA Graph Optimization for Low-Latency Large Language Model Inference

Resumo Técnico: Otimização Híbrida JIT–CUDA Graph para Inferência de LLMs de Baixa Latência

1. O Problema (Contexto e Motivação)

A implantação de Grandes Modelos de Linguagem (LLMs) em aplicações interativas (como assistentes de código ou chatbots) enfrenta um gargalo crítico: a latência de inferência. Durante o processo de decodificação autorregressiva (token por token), o sistema sofre com um alto custo de overhead causado pelo lançamento repetitivo de kernels na GPU, eventos de sincronização e a coordenação entre a CPU (host) e a GPU (device).

Atualmente, existe um dilema técnico:

• CUDA Graphs: Oferecem execução extremamente rápida ao pré-capturar e "reproduzir" sequências de kernels, mas exigem formas de tensores (shapes) estáticas e fluxo de controle determinístico, o que é incompatível com a natureza dinâmica dos LLMs (como amostragem estocástica e comprimentos de sequência variáveis).
• Compilação JIT (Just-In-Time): Oferece a flexibilidade necessária para lidar com formas dinâmicas e fluxos de controle, mas ainda incorre em latência devido ao despacho de kernels via CPU a cada passo de decodificação.

2. Metodologia (A Abordagem Híbrida)

Os autores propõem um framework de runtime híbrido que particiona a inferência do Transformer em dois domínios distintos, operando de forma assíncrona:

• Domínio Estático (CUDA Graph Replay): Componentes computacionalmente intensos, mas com formas de tensores previsíveis (como projeções lineares, normalização de camada e cálculos de atenção), são capturados em CUDA Graphs. Isso permite que a GPU execute o grafo de computação sem intervenção constante da CPU, reduzindo drasticamente o tempo de lançamento de kernels.
• Domínio Dinâmico (JIT Compilation): Componentes que dependem de condições de tempo de execução (como amostragem de tokens, atualizações do cache KV e extensões de embeddings posicionais) são gerenciados via compilação JIT (utilizando a infraestrutura do PyTorch). Isso preserva a flexibilidade necessária para a decodificação autorregressiva.

Inovação de Implementação: O sistema utiliza um mecanismo de geração assíncrona de grafos. Enquanto o motor de reprodução (Replay) executa um grafo já existente, um processo em segundo plano captura novos grafos para diferentes comprimentos de sequência, armazenando-os em um buffer rotativo. Isso permite que o sistema "aprenda" e se adapte a novos comprimentos de contexto sem bloquear a execução atual.

3. Principais Contribuições

• Particionamento Inteligente: Divisão clara entre operações estáticas (seguras para grafos) e dinâmicas (que exigem JIT).
• Captura Assíncrona: Implementação de um mecanismo que sobrepõe a captura de novos grafos com a reprodução de grafos existentes, minimizando o tempo de espera.
• Redução de Variância: O uso de grafos reduz o jitter (oscilação) de latência, tornando o tempo de resposta mais previsível.
• Arquitetura Desacoplada: Coordenação via comunicação entre processos (IPC) para separar a geração de contexto da geração de grafos.

4. Resultados Experimentais

Os testes foram realizados utilizando o modelo LLaMA-2 7B em uma GPU NVIDIA H100, comparando a abordagem proposta com o PyTorch Eager (HuggingFace) e o TensorRT–LLM.

• Time-to-First-Token (TTFT): O framework híbrido reduziu o TTFT em até 66% em comparação com o PyTorch Eager e superou consistentemente o TensorRT–LLM em cenários de sequências curtas a médias.
• Latência de Cauda (P99): O método alcançou uma redução de 20,2% na latência P99 em relação ao TensorRT–LLM, o que significa que o sistema é muito mais estável e sofre menos com atrasos inesperados em tokens individuais.
• Escalabilidade: O tempo de resposta escalou de forma quase linear com o comprimento do contexto, demonstrando maior estabilidade que os modelos de base.

5. Significância e Conclusão

Este trabalho é significativo porque resolve o compromisso (trade-off) entre eficiência estática e flexibilidade dinâmica. Ao contrário de soluções que exigem a reescrita de operadores em C++ ou que assumem formas de tensores fixas, a abordagem híbrida é aplicável a modelos autorregressivos reais e dinâmicos.

A pesquisa demonstra que, para aplicações de IA sensíveis à latência (onde o usuário espera uma resposta imediata e fluida), a combinação de compilação JIT com a reprodução de CUDA Graphs é uma estratégia de otimização superior para o regime de decodificação de curto e médio alcance.