Skip to content

Spec: ML-KEM NTT and inverse of a polynomial stored with 16 bits - #3042

Merged
alex merged 1 commit into
mainfrom
spec-mlkem16-ntt
Oct 21, 2026
Merged

alex merged 1 commit into
mainfrom
spec-mlkem16-ntt

Conversation

@alex

@alex alex commented Oct 21, 2026

Copy link
Copy Markdown
Member

Trust change: new Spec (lean/VerifiedGarbage/Spec/MlKem/Ntt16.lean). It adds the contracts and Apis of two new functions:

  • vg_mlkem16_ntt(f: *mut [u16; 256], scratch: *mut [u64; 64]): NTT (FIPS 203 Algorithm 9) in place.
  • vg_mlkem16_inv_ntt(f: *mut [u16; 256], scratch: *mut [u64; 64]): NTT⁻¹ (Algorithm 10) in place.

They are the 16-bit counterparts of vg_mlkem_ntt and vg_mlkem_inv_ntt.

Why

The AArch64 ML-KEM tops inline the 16-bit NTT today, in 4 to 11 copies per function. The next PR replaces it with a rescheduled transform, in the style of mlkem-native: by-element zetas, fewer transposes in layers 4–7, and two blocks interleaved per trip. That doubles the transform's size, so the tops will call it out of line instead of inlining it. A call needs a registered function with a contract, and none exists for 16-bit storage: vg_mlkem_ntt takes u32 coefficients. The inverse is included so that its rescheduling can follow without another Spec PR.

The contract (inPlace16Contract)

  • Pre: the polynomial at f is reduced with 16-bit coefficients (ReducedW 16).
  • Post: f holds t of it, reduced (PolyIsW 16 … (t (polyAtW 16 m f))), with t = ntt or t = nttInv.
  • Scratch: working space of 512 bytes, which the caller destroys.
  • Unchanged definitions: the storage predicates are those of Spec/MlKem/Kpke16.lean, and ntt and nttInv are those of Spec/MlKem/Poly.lean.
  • Table of zetas: an implementation's table is a registration consts entry, as today, not an argument.

Nothing implements or calls these yet, so no code is emitted and src/asm/ is unchanged.

Checks

  • lake build of the module.
  • The Python checks:
    • check_lean_imports, check_lean_modules and check_lean_public_imports;
    • check_lean_speed and check_registrations.

Generated by Claude Code

@alex
alex added this pull request to the merge queue Oct 21, 2026
Merged via the queue into main with commit 596fe1b Oct 21, 2026
11 checks passed
@alex
alex deleted the spec-mlkem16-ntt branch October 21, 2026 21:09
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Development

Successfully merging this pull request may close these issues.

2 participants