Verifying AWS Neuron
NKI Kernels with ESBMC

Context

NKI is the Python-level kernel API that AWS exposes for writing high-performance tensor kernels on Trainium and Inferentia NeuronCores. NKI kernels look like ordinary Python but are subject to hardware-specific shape, dtype, and partition constraints, for example, the partition dimension PMAX is fixed at 128, and the matmul unit has separate stationary/moving free-dimension limits (GEMM_STATIONARY_FMAX = 128, GEMM_MOVING_FMAX = 512). Violating any of these constraints produces a hardware-level error that is painful to diagnose at runtime.

The question we wanted to answer is simple: can these constraints be checked formally and exhaustively, before deployment, on a developer laptop and without a NeuronCore? If so, NKI authors get a fast, deterministic safety net that complements on-device testing rather than competing with it.

Approach

We wrote a single ~865-line stubs.py library modelling NKI tile shapes, dtypes and hardware constraints, so that upstream NKI kernels could run unchanged on top of the stubs. Ported kernels are byte-for-byte faithful to the upstream sources for indexing; no source rewriting is required. The harness is pinned to the aws-neuron/nki-samples snapshot at commit a87aaa44, covering 11 source files (6 tutorials and 5 community kernels) and 19 kernel functions.

Verification splits into two phases. Phase 1 checks shapes and bounds via the contract assertions in stubs.py: partition-dimension limits, in-bounds slicing, DMA / elementwise / matmul shape-equality, and matmul-unit hardware bounds. Phase 2 turns on ESBMC's --overflow-check (with default division-by-zero detection) to catch CWE-190 (signed integer overflow) and CWE-369 (division by zero) in the host-side index arithmetic that surrounds the kernels. Both phases use ESBMC 8.3.0 or later with the default Bitwuzla SMT solver.

For symbolic-shape targets, where input dimensions are non-deterministic, we used k-induction to obtain completeness certificates: 11 of 13 symbolic targets establish that the chosen --unwind bound is sufficient to cover every reachable state, so the verdict is exhaustive, not merely bounded. The two exceptions (the Mamba-v3 and attention-forward-v3 symbolic targets) timed out under k-induction and remain heuristic; we report this openly in the repository.

Results

What this means

NKI shape and bounds checking is amenable to bounded model checking with k-induction completeness, at developer-laptop speeds, with no hardware in the loop. The same technique extends to integer-overflow and division-by-zero defects in the host-side glue around kernels. For teams writing or reviewing NKI code, this turns a class of hardware-level failures into compile-time CI gates, and the cost of adoption is a stub library plus a handful of minutes per kernel.