Verificando Kernels NKI
do AWS Neuron com ESBMC

Contexto

O NKI é a API de kernel em nível de Python que a AWS expõe para escrever kernels de tensor de alto desempenho nos NeuronCores Trainium e Inferentia. Kernels NKI parecem Python comum, mas estão sujeitos a restrições de forma, dtype e partição específicas do hardware, por exemplo, a dimensão de partição PMAX é fixada em 128, e a unidade de matmul tem limites de free-dimension stationary/moving separados (GEMM_STATIONARY_FMAX = 128, GEMM_MOVING_FMAX = 512). Violar qualquer uma dessas restrições produz um erro em nível de hardware doloroso de diagnosticar em tempo de execução.

A pergunta que quisemos responder é simples: essas restrições podem ser verificadas formal e exaustivamente, antes da implantação, em um laptop de desenvolvedor e sem um NeuronCore? Se sim, autores de NKI ganham uma rede de segurança rápida e determinística que complementa os testes em dispositivo em vez de competir com eles.

Abordagem

Escrevemos uma única biblioteca stubs.py de ~865 linhas modelando formas de tiles, dtypes e restrições de hardware do NKI, de modo que kernels NKI upstream pudessem rodar sem alterações sobre os stubs. Os kernels portados são fiéis byte a byte às fontes upstream quanto à indexação; nenhuma reescrita de código-fonte é necessária. O harness está fixado ao aws-neuron/nki-samples no snapshot do commit a87aaa44, cobrindo 11 arquivos-fonte (6 tutoriais e 5 kernels da comunidade) e 19 funções de kernel.

A verificação se divide em duas fases. A Fase 1 verifica formas e limites por meio das asserções de contrato em stubs.py: limites de dimensão de partição, fatiamento dentro dos limites, igualdade de forma de DMA / elementwise / matmul e limites de hardware da unidade de matmul. A Fase 2 ativa o --overflow-check do ESBMC (com detecção padrão de divisão por zero) para capturar CWE-190 (overflow de inteiro com sinal) e CWE-369 (divisão por zero) na aritmética de índices do lado do host que envolve os kernels. Ambas as fases usam o ESBMC 8.3.0 ou posterior com o solver SMT Bitwuzla padrão.

Para alvos de forma simbólica, onde as dimensões de entrada são não-determinísticas, usamos k-induction para obter certificados de completude: 11 de 13 alvos simbólicos estabelecem que o limite de --unwind escolhido é suficiente para cobrir todo estado alcançável, de modo que o veredito é exaustivo, não meramente limitado. As duas exceções (os alvos simbólicos Mamba-v3 e attention-forward-v3) estouraram o tempo limite sob k-induction e permanecem heurísticas; reportamos isso abertamente no repositório.

Resultados

O que isso significa

A verificação de formas e limites do NKI é tratável por verificação limitada de modelos com completude por k-induction, em velocidades de laptop de desenvolvedor, sem hardware no loop. A mesma técnica se estende a defeitos de overflow de inteiro e divisão por zero no glue do lado do host em torno dos kernels. Para equipes que escrevem ou revisam código NKI, isso transforma uma classe de falhas em nível de hardware em gates de CI em tempo de compilação, e o custo de adoção é uma biblioteca de stubs mais alguns minutos por kernel.