Writing NKI Kernels

This skill guides writing and modifying NKI (Neuron Kernel Interface) kernels — from new kernel creation (PyTorch/NumPy/natural language translation) to editing existing kernels (adding shape/dtype support, refactoring tiling, implementing new features). Focus on correctness using documented APIs.

Critical: NKI Language Constraints

BEFORE writing any NKI code, read references/nki-language-constraint.md for the complete list of required and forbidden API patterns covering Beta 1 → Beta 2, Beta 2 → NKI 0.3.0, and NKI 0.3.0 → NKI 0.4.0 migration rules. Violating ANY rule is a compilation failure.

Hardware limits: don't hardcode the PSUM free dim

The PSUM free-dimension limit is generation- and dtype-gated and has changed across releases. Current values:

• gen2 / gen3: 512 (one PSUM bank)
• gen4: 4096 for a float32 dst, 8192 for a bfloat16 dst (entire PSUM)