Equivalence Checking of ML GPU Kernels

Abstract

With the rapid progress of deep learning and large language models (LLMs), companies spend enormous sums executing GPU kernels. These kernels have become prime targets for aggressive optimization. Recent efforts increasingly leverage LLMs to generate GPU kernels, but make no formal guarantees about the generated kernels. We present the first equivalence checker for GPU kernels and use it to formally verify the correctness of machine learning (ML) kernels optimized by hand, by LLM, and by compiler. We show that our equivalence checker is sound and, for a well-defined class of GPU kernels which includes many programs of interest, complete. Our implementation, Volta, can verify ML computations such as convolutions, matrix multiplications, and various attention mechanisms.