NKIBench is a benchmark of AWS Neuron Kernel Interface (NKI) kernels paired with NumPy reference implementations. Each task provides a specification, a ground-truth NumPy forward pass, and an optimized NKI kernel targeting AWS Trainium devices, together with tooling to compile, check numerical correctness, and measure on-device latency. Paper.
NKIBench/
├── seeds/ # YAML task specifications (shape-agnostic templates)
├──… See the full description on the dataset page:
https://huggingface.co/datasets/Genghan/NKIBench.
