Static Verification of AWS Neuron NKI Kernels Using ESBMC: Catching Shape and Boundary Errors Before Compilation

This article introduces a project developed by lucasccordeiro that uses the ESBMC model checker for static verification of AWS Neuron NKI kernels, which can catch shape and boundary-related errors before compilation or runtime. This method combines a Python stub library with symbolic execution to help developers shorten debugging cycles, reduce hardware dependencies, and increase code confidence.

AWS Neuron is a deep learning SDK for Inferentia and Trainium accelerators, and NKI allows developers to write custom kernels. However, currently NKI users can only detect shape and boundary errors during compilation or runtime, leading to longer development cycles and high debugging costs.

The project provides a lightweight Python stub library to simulate NKI's tile shape and boundary checks, allowing ESBMC to detect issues at the symbolic execution level. The stub library (stubs.py) tracks shapes and dtypes, covering ndarray allocation, slicing, data transfer, ISA operations, etc. The build process concatenates stubs, kernels, and harness files into a single file (since the ESBMC Python frontend cannot resolve cross-module imports). Each target verification takes 1-3 seconds, and the complete suite finishes within 15 seconds.

ESBMC can verify the following key preconditions:

• Partition dimension constraints (SBUF tile allocation ≤128)
• Boundary validity of 2D tile slices in loop iterations
• Input-output shape consistency for operations like nisa.dma_copy
• Data type consistency for the above operations
• Divisibility/Tile count assertions at kernel entry
• Shape constraints on kernel return values

Target	Type	Expected Result	Actual Result
tensor_add	Correct Implementation	SUCCESSFUL	✅ Passed
tensor_add_buggy	Intentionally Buggy	FAILED	✅ Detected
tensor_add_symbolic	Symbolic Shape	SUCCESSFUL	✅ Passed
transpose2d	Correct Implementation	SUCCESSFUL	✅ Passed
transpose2d_buggy	Intentionally Buggy	FAILED	✅ Detected
matmul	Correct Implementation	SUCCESSFUL	✅ Passed
matmul_big	Large Size	SUCCESSFUL	✅ Passed
matmul_buggy	Intentionally Buggy	FAILED	✅ Detected
All cases met expectations: correct implementations passed verification, and intentionally buggy ones were successfully detected.

Current Limitations:

• Does not verify numerical correctness (no involvement with ISA semantics or numerical computation)
• No SPMD interaction verification (does not cover multi-core parallelism)
• Does not support complex indexing (e.g., broadcast indexing for maxpooling, interpolate)
• Symbolic variants require manual setting of unwind boundaries (k-induction not integrated)

Future Extensions: Need to add about a dozen stub functions to cover the nki-samples corpus, including reductions (nl.sum), elementwise operations (nisa.tensor_scalar), access patterns (.ap()), broadcast indexers (nl.ds, par_dim), and decorators (@nki.jit, etc.)

Value for NKI developers:

1. Shorten debugging cycles: Catch slice arithmetic errors before submitting to the Neuron compiler
2. Reduce hardware dependencies: Verify without real Trainium/Inferentia devices
3. Increase code confidence: Formal verification provides guarantees far beyond unit test coverage

This project demonstrates the practical value of formal methods in AI accelerator programming. By combining ESBMC symbolic execution with the NKI Python frontend, errors that would only be exposed at runtime can be caught before compilation. As the NKI ecosystem develops, this static verification method is expected to become a standard practice in kernel development.