Status: v0.1 — 3 recipes (popcount / xor-popcount / hamming-distance)
- tutorials + benchmarks + honest failure transcripts.
A community-curated cookbook of practical kernel recipes for AWS Neuron's NKI (Neuron Kernel Interface). Each recipe is a complete, runnable, benchmarked example showing how to implement a building-block operation on Inferentia2 / Trainium hardware — and where it doesn't.
nki-cookbook is not an official AWS product and is not affiliated with,
endorsed by, or sponsored by Amazon Web Services. NKI itself is an AWS
product; this repository is a third-party collection of recipes built on
top of it. The names "AWS", "Neuron", "NeuronX", "Inferentia", and
"Trainium" are trademarks of Amazon Web Services, used here for
nominative identification only.
The official nki-samples is
an excellent reference. It contains:
- High-performance reference kernels (FlashAttention, etc.) for ML workloads
- Tutorial kernels mirroring the NKI guides
nki-cookbook is complementary, not competitive. It focuses on:
- General-purpose primitives outside of standard ML (bit ops, compression, set algebra) for use in databases, search engines, analytics systems
- Practical recipes with self-contained README + benchmark + results, designed to be copy-pasted into application code
- Honest comparisons against NVIDIA L4 GPU and host CPU baselines on the same workload — including cases where Inferentia loses
- Failure transcripts (
failures/) documenting patterns that look like they should work but don't, including[XCG863] ISA check failedfor the obvious bit-packed Hamming kernel on NCv2
If you want ML primitives, use nki-samples. If you want bit-twiddling,
compression, and an honest answer to "should I use Inferentia for this?",
you're in the right place.
Every recipe in recipes/ follows the same four-file layout:
recipes/01-popcount/
├── README.md # Problem, approach, perf characteristics, gotchas
├── implementation.py # NKI kernel + nki.simulate_kernel test
├── benchmark.py # Reproducible benchmark harness
└── results.md # Measured numbers (Inf2 + L4 + CPU baseline)
This means you can:
- Read the README to decide if the recipe is right for your use case
- Copy
implementation.pyinto your project (Apache-2.0) - Run
benchmark.pyon your own hardware to verify the numbers - Trust
results.mdbecause it was generated bybenchmark.pyon real Inf2 / L4 hardware (provenance recorded inline)
| # | Recipe | Use case | Approach | At n=131,072 |
|---|---|---|---|---|
| 01 | popcount | bit counting (BBQ, bloom filters, set cardinality) | NKI shift+mask+add OR uint32 SWAR | educational; CPU AVX-LUT wins for popcount-only |
| 02 | xor-popcount | building block of Hamming distance | unpacked nki.language.bitwise_xor + reduce-sum |
memory-bound by 8× unpack |
| 03 | hamming-distance | BBQ vector search, bit vectors comparison | matmul reformulation via Tensor Engine (BF16/INT8 dot product) | 11.77M pairs/s on inf2.xlarge |
NKI simulator parity vs NumPy is bit-exact for all three recipes.
| Recipe primitive | NCv2 (Inf2) | NCv3+ (Trn2) |
|---|---|---|
nki.language.bitwise_xor |
✅ supported | ✅ supported |
nki.isa.nonzero_with_count |
❌ not available | ✅ supported |
uint64 dtype in NKI ops |
❌ not available | (verify when v0.2 lands) |
uint32 dtype |
✅ supported | ✅ supported |
| 128-row partition tiling | required (axis cap = 128) | required |
| Packed XOR + popcount-loop kernel | ❌ [XCG863] ISA check fail (failure) |
(verify) |
These limitations are documented in each recipe's README.md under
"Hardware compatibility". If you're targeting Inf2, use the unpacked +
matmul path; if you're targeting Trn2, you can switch to the
nonzero_with_count fast path (and probably get more upside from
INT8 Tensor Engine).
The headline finding from recipe 03's bench harness, on a 131,072-document corpus of 768-bit BBQ vectors:
| Hardware | BF16 matmul pairs/s | Notes |
|---|---|---|
inf2.xlarge (1 NeuronCore-v2) |
11.77M | Tensor Engine + HBM 32 GiB |
g6.xlarge (NVIDIA L4) |
10.42M | Tensor Cores + GDDR6 24 GiB |
| inf2.xlarge host CPU (4 vCPU) | 3.64M | NumPy + AVX-LUT popcount |
L4 is competitive at all batch sizes and wins for n ≤ 8,192. Inf2 only
edges ahead at very large N (~13% advantage at n=131k saturation). The
real Inf2 advantage is $/query at spot prices (~2.8× cheaper than L4
spot in us-west-2 as of 2026-05-10) and HBM capacity (32 GiB vs
24 GiB). See benchmarks/decision-tree.md
for when to pick which.
# Clone
git clone https://github.com/abyo-software/nki-cookbook
cd nki-cookbook
# Install AWS Neuron SDK 2.24+
# https://awsdocs-neuron.readthedocs-hosted.com/
pip install --extra-index-url https://pip.repos.neuron.amazonaws.com \
"neuronx-cc==2.24.*" torch torch-neuronx
# Run a recipe in CPU simulator (no Inf2 needed)
# All recipes verified against NumPy reference via nki.simulate_kernel
cd recipes/01-popcount
python implementation.py
# Run benchmark on real Inf2 (requires inf2.* instance)
python benchmark.py- bit packing / unpacking (SIMDBP128 family) — Wave 2
- delta encoding — Wave 2
- Roaring Bitmap operations (intersection, union, difference) — Wave 2
- simple PFOR compression — Wave 2
- prefix sum / scan — Wave 2 / 3
- top-k selection — Wave 3
Wave timing is tied to the companion projects below.
nki-cookbook is the educational / honest-comparison output of broader
work on Inferentia2 search engine extensions:
- Phase 0.5 (BBQ Hamming) verdict and recipes 01-03 source — see this
cookbook's
failures/andbenchmarks/ - Postings list / doc_values integer compression on Inf2 — recipes 04-08 source (planned for v0.2)
- ML node features (anomaly detection, classification, forecast, etc.) on Optimum Neuron — recipes 09+ source (planned for v0.3)
tutorials/getting-started.md— local development withnki.simulate_kernel, no Inf2 hardware requiredtutorials/optimization-guide.md— 9 NKI design patterns (works / doesn't work / pitfalls)
benchmarks/compare-with-cpu-gpu.md— full Inf2 vs L4 vs CPU sweep across batch sizes and dtype combinationsbenchmarks/decision-tree.md— "should I use Inferentia for this?" decision tree based on workload shape
failures/packed-xor-popcount-NCv2.md—[XCG863] ISA check failedfor the obvious bit-packed Hamming kernelfailures/at-nki-jit-recompile.md—@nki.jitbaremetal mode pays ~8 s recompile per call
This is a best-effort project. We aim to:
- Acknowledge new-recipe proposals within 7 days
- Review benchmark reproducibility within 30 days
- Cut versioned releases when meaningful new recipes land
We do not guarantee:
- Coverage of every NeuronX feature
- Production support
- Backward compatibility before v1.0
See CONTRIBUTING.md. New recipes welcome — open an issue describing the problem before coding.
Apache License 2.0 — see LICENSE.