Model2Kernel: Model-Aware Symbolic Execution For Safe CUDA Kernels

O artigo apresenta o Model2Kernel, o primeiro sistema prático para verificar automaticamente a segurança de memória de kernels CUDA usados em inferência de LLMs, combinando análise dinâmica orientada ao modelo com execução simbólica especializada para detectar bugs que métodos existentes não conseguem identificar.

O essencial

Imagine que você tem um restaurante de luxo (o sistema de Inteligência Artificial) que serve milhões de pratos complexos todos os dias. Para cozinhar esses pratos com velocidade incrível, o restaurante usa uma equipe de milhares de chefs robóticos trabalhando em perfeita sincronia. Esses chefs são os Kernels CUDA (pequenos programas que rodam na placa de vídeo/GPU).

O problema é que, como esses chefs trabalham tão rápido e lidam com ingredientes que mudam de tamanho a cada pedido (o tamanho da conversa do usuário), eles às vezes cometem erros graves:

• Tentam pegar um ingrediente que não existe na geladeira (Acesso fora dos limites).
• Usam uma receita que exige um número de ovos maior do que o que cabe no balcão de números inteiros (Estouro de inteiro).
• Dois chefs tentam usar a mesma faca ao mesmo tempo, causando uma briga (Corrida de dados).

Esses erros podem fazer o restaurante inteiro fechar (o sistema travar) ou, pior, permitir que um ladrão entre na cozinha e roube as receitas secretas (os pesos do modelo).

O Problema: Como encontrar esses erros?

Até agora, tentar achar esses erros era como tentar encontrar um fio de cabelo em um palheiro:

1. Testar na vida real (Execução Dinâmica): Você precisava de um restaurante gigante e caro apenas para testar. Se o erro só acontece quando o cliente pede um prato gigante, você precisaria esperar esse cliente aparecer.
2. Ler as receitas (Análise Estática): Tentar ler o código manualmente ou com ferramentas antigas era impossível porque as receitas mudavam de tamanho dependendo do cliente. As ferramentas antigas diziam: "Não consigo analisar isso, o tamanho do ingrediente é desconhecido".

A Solução: O "Model2Kernel"

Os autores criaram o Model2Kernel, que é como um super-inspetor de cozinha que não precisa de ingredientes reais para descobrir onde estão os erros. Ele funciona em duas etapas principais, como se fosse uma dupla de detetives:

1. O "HFProbe": O Espião que Entende o Cardápio

Primeiro, o sistema precisa entender como o restaurante funciona. O HFProbe é um espião que observa o modelo de IA (o cardápio) sem precisar ligar os fogões reais.

• Ele olha para a receita e diz: "Ok, para este prato, o tamanho do ingrediente 'batch' (lote) é fixo pela receita, mas o tamanho do 'pedido' (sequência de texto) é definido pelo cliente."
• Ele também muda o cardápio propositalmente (mutação) para ver se consegue forçar os chefs a tentarem receitas que eles normalmente não fariam, garantindo que nenhum canto da cozinha seja deixado de fora.

2. O "cuKLEE": O Simulador de Realidade Paralela

Agora que sabemos quais ingredientes são fixos e quais variam, o cuKLEE entra em ação. Ele é um simulador de realidade paralela especializado em cozinhas de robôs.

• Em vez de cozinhar um prato de cada vez, ele cria uma "bolha de pensamento" onde os ingredientes podem ter qualquer tamanho possível ao mesmo tempo.
• Ele usa "ingredientes simbólicos" (como variáveis mágicas) para representar o tamanho do pedido.
• Ele pergunta ao cérebro lógico (o solucionador de restrições): "Existe algum tamanho de pedido onde o chef vai tentar pegar um ingrediente que não existe?"
• Se a resposta for "Sim", ele aponta exatamente onde o erro acontece e diz: "Atenção! Se o cliente pedir 128.000 palavras, o chef vai quebrar a mesa na linha X!"

Por que isso é genial?

• Não precisa de hardware caro: Você não precisa de uma placa de vídeo gigante para testar. O sistema "pensa" o suficiente para saber onde vai dar errado.
• Entende a complexidade: Ele sabe que os chefs trabalham em equipe (milhares de threads) e verifica se eles não vão brigar entre si.
• Resultados Reais: Ao testar em modelos famosos (como os do vLLM e Hugging Face), o sistema encontrou 353 erros que ninguém tinha visto antes! A maioria eram erros de cálculo de tamanho (inteiros estourando) que levariam a falhas catastróficas.

A Analogia Final

Imagine que você está construindo uma ponte para carros que podem ter tamanhos variados (de uma bicicleta a um caminhão de 50 toneladas).

• Os métodos antigos diziam: "Vamos construir a ponte e esperar que um caminhão de 50 toneladas passe para ver se ela cai." (Perigoso e caro).
• O Model2Kernel diz: "Vamos usar matemática avançada para simular todos os tamanhos de veículos possíveis ao mesmo tempo. Vamos descobrir que, se um caminhão de 40 toneladas passar na curva X, a ponte vai quebrar, e vamos consertar isso antes de colocar o primeiro tijolo."

Conclusão

O Model2Kernel é uma ferramenta de segurança que garante que a infraestrutura que roda nossas IAs (como chatbots e assistentes) seja sólida. Ele previne que pequenos erros de cálculo transformem um sistema inteligente em um sistema quebrado ou vulnerável a hackers, tudo isso de forma automática e sem precisar de equipamentos caros. É como ter um guarda-costas que prevê ataques antes mesmo deles acontecerem.

Resumo técnico

1. O Problema

A adoção generalizada de Grandes Modelos de Linguagem (LLMs) tornou a inferência acelerada por GPU um componente crítico da infraestrutura de computação moderna. Sistemas de inferência de produção dependem de kernels CUDA para implementar operações centrais de transformadores. No entanto, esses kernels são altamente suscetíveis a erros de segurança de memória devido a três fatores principais:

1. Layouts de Tensores Dependentes do Modelo: A estrutura de memória dos tensores varia conforme a arquitetura do modelo e o tamanho da sequência de entrada, tornando difícil prever limites de memória estáticos.
2. Indexação de Memória Intrincada: Cálculos de índice de grão fino em grandes tensores e milhares de threads aumentam o risco de acessos fora dos limites (out-of-bounds) e condições de corrida (data races).
3. Paralelismo Massivo e Interações Complexas: O crescimento das arquiteturas de modelos e o aumento do comprimento das sequências de inferência elevam o risco de estouro de inteiros (integer overflows), que podem levar a acessos de memória inválidos.

Consequências: Esses erros podem corromper pesos do modelo, causar falhas catastróficas no serviço de inferência ("crash") ou, mais grave, permitir ataques adversariais, como a execução de código arbitrário na GPU ou a modificação de pesos implantados.

Limitações das Técnicas Existentes:

• Métodos Dinâmicos: Dependem de hardware não disponível ou incorrem em sobrecarga de tempo de execução muito alta.
• Ferramentas Estáticas: Não conseguem lidar com tamanhos de buffer determinados em tempo de execução ou com a variabilidade das configurações de lançamento de kernels.
• Falta de Abordagem Específica: Nenhuma técnica existente consegue detectar efetivamente erros de memória em kernels CUDA dentro de sistemas de inferência de LLMs modernos.

---

2. Metodologia: Model2Kernel

O Model2Kernel é o primeiro sistema prático para verificar automaticamente a segurança de memória de kernels CUDA usados em inferência de LLMs. Ele combina análise dinâmica orientada ao modelo com execução simbólica especializada. O sistema é composto por dois componentes principais:

A. HFProbe (Perfilador de Modelos)

O HFProbe executa o modelo (sem necessidade de hardware GPU real) para entender como o modelo invoca os kernels CUDA.

• Função: Identifica quais argumentos do kernel são fixos (determinados pela arquitetura do modelo) e quais são controlados pelo usuário (variáveis).
• Técnicas:
  • Realiza análise estática do código Python do modelo para construir um grafo de chamadas.
  • Executa o modelo com configurações reais usando kernels CUDA "falsos" (mocked) para registrar como os tensores são manipulados.
  • Mutação de Configuração: Utiliza um agente de LLM para modificar arquivos de configuração (config.json) e parâmetros do framework, forçando a execução de kernels que não seriam ativados por configurações padrão, garantindo uma cobertura mais ampla.
• Saída: Gera um contexto de análise com restrições concretas sobre as formas dos tensores e valores de parâmetros, reduzindo o espaço de busca para o executor simbólico.

B. cuKLEE (Motor de Execução Simbólica para CUDA)

O cuKLEE é um motor de execução simbólica especializado em CUDA, baseado no KLEE, mas adaptado para lidar com as particularidades de inferência de LLM.

• Modelo de Memória de Tensores: Em vez de tratar a memória como um bloco contíguo, o cuKLEE modela cada tensor como uma região de memória distinta e separada. Isso simplifica a verificação de limites, pois um ponteiro de um tensor nunca pode referenciar outro tensor através de aritmética de ponteiros.
• Variáveis Simbólicas:
  • Tensores: Representa endereços base, contagem de elementos, dimensões e tipos como variáveis simbólicas.
  • Threads: Trata identificadores de thread (threadIdx.x, blockIdx.x) e blocos como variáveis simbólicas, permitindo analisar milhares de threads em uma única passagem de execução simbólica, em vez de enumerá-los.
• Resumo de Métodos de Tensor: Agrupa e resume 140 métodos de tensores do PyTorch (como numel(), sum(), size()) para evitar a análise profunda de lógica complexa de gerenciamento de memória, mantendo a precisão das restrições.
• Detecção de Bugs: Verifica quatro tipos de vulnerabilidades:
  1. Acessos fora dos limites (Out-of-bounds).
  2. Estouro de inteiros (Integer overflows).
  3. Condições de corrida (Data races).
  4. Desreferenciamento de ponteiros nulos.

---

3. Contribuições Principais

1. cuKLEE: Um motor de execução simbólica específico para CUDA que suporta tensores com formas dinâmicas e escala para milhares de threads de CUDA através de abstrações de memória e identificadores de thread simbólicos.
2. HFProbe: Um perfilador dinâmico de modelos que analisa como os modelos de LLM invocam kernels CUDA, inferindo argumentos dependentes do modelo e gerando contextos de análise otimizados.
3. Model2Kernel: Um framework totalmente automatizado que integra o cuKLEE e o HFProbe para detectar bugs de memória em sistemas de inferência de LLM.
4. Validação Empírica: Demonstração da eficácia do sistema através da descoberta de centenas de bugs em modelos e kernels reais.

---

4. Resultados Experimentais

O sistema foi avaliado em kernels e modelos de três fontes: vLLM, Hugging Face e artigos de pesquisa recentes (4 papers).

• Descoberta de Bugs:
  • Total: 353 bugs de memória previamente desconhecidos foram encontrados.
  • Detalhamento: 328 estouro de inteiros (integer overflows) e 25 acessos fora dos limites (out-of-bounds).
  • Falsos Positivos: Apenas 9 (taxa de 2,49%), demonstrando alta precisão.
• Comparação com o Estado da Arte:
  • Ao testar 20 bugs conhecidos do vLLM, o Model2Kernel (via cuKLEE) detectou 15, enquanto as ferramentas de base (Honeycomb, GKLEE, ESBMC-GPU) detectaram no máximo 5.
  • As ferramentas existentes falharam principalmente devido à incapacidade de lidar com tamanhos de memória dinâmicos, configurações de lançamento variáveis e complexidade do código C++ moderno.
• Custo: A análise total levou cerca de 986 horas (agregadas em processos paralelos) e custou aproximadamente $442 em uso de API de LLM para mutação de configuração, o que é considerado moderado dado o tamanho do código analisado (>100k linhas).

Estudo de Ablação:

• Sem o HFProbe, o número de falsos positivos saltou para mais de 5.000, pois o executor simbólico inferiu formas de tensores teoricamente possíveis, mas impossíveis na prática.
• Sem a mutação de configuração, 32 bugs adicionais foram perdidos, pois muitos kernels não eram ativados pelas configurações padrão.

---

5. Significado e Impacto

O Model2Kernel representa um avanço significativo na segurança de sistemas de IA:

• Segurança Crítica: Ao detectar vulnerabilidades de memória em kernels CUDA, o sistema previne falhas de serviço e, crucialmente, mitiga riscos de segurança onde adversários poderiam manipular pesos de modelos ou executar código malicioso na GPU.
• Viabilidade Prática: Diferente de abordagens puramente estáticas (que falham em cenários dinâmicos) ou puramente dinâmicas (que são lentas ou exigem hardware), o Model2Kernel oferece uma solução automatizada e escalável para a indústria de LLMs.
• Generalização: As ideias centrais do trabalho (modelagem de tensores como regiões discretas, simbolização de IDs de threads e perfilagem de modelos) podem ser estendidas para outros tipos de bugs e até para kernels em NPUs (Unidades de Processamento Neural).

Em resumo, o Model2Kernel preenche uma lacuna crítica na infraestrutura de inferência de LLMs, fornecendo a primeira ferramenta prática para garantir a segurança de memória em kernels CUDA complexos e dinâmicos.