[alex]

Provenance. The forward NTT's schedule is ported from mlkem-native's mlkem_ntt_aarch64_asm.S, its dev/aarch64_opt version scheduled by SLOTHY. The source is the copy vendored in AWS-LC: aws-lc-sys 0.45.0, aws-lc/crypto/fipsmodule/ml_kem/mlkem/native/aarch64/src/mlkem_ntt_aarch64_asm.S.

• Copyright © 2022 Arm Limited, © 2022 Hanno Becker, © 2023 Amin Abdulrahman and Matthias Kannwischer, © the mlkem-native project authors.
• License: Apache-2.0 OR ISC OR MIT. The header of Impl/MlKem/AArch64/Ntt16/Fast.lean keeps the notice.
• Its instructions are transcribed into Lean Instr lists, then proven like every other function here. Nothing is trusted because it came from mlkem-native.

It builds on #3020 (TCB: the by-element multiplies), #3042 (Spec: vg_mlkem16_ntt's Api) and #3041 (the NTT machine), all merged.

What changes

• vg_mlkem16_ntt is a new out-of-line function on AArch64: Impl/MlKem/AArch64/Ntt16/Fast.lean.
  • Layers 1–3 run a column of eight vectors at a time, with the zetas multiplied by element (MUL/SQRDMULH/MLS by element).
  • Layers 4–7 run a 32-coefficient block at a time, with fewer transposes.
  • Both loops are software-pipelined as SLOTHY scheduled them.
• Changes from the original:
  • Our zetas are in [0, q), so that the bounds of Spec's transform hold. Each is stored with its tw = round(ζ·2¹⁵/q), as mlkem-native's are.
  • The two tables are one constant, VG_MLKEM16_NTT_ZETAS (928 bytes); the second starts 160 bytes in.
  • A final loop reduces each coefficient to [0, q), which ntt16Contract requires.
  • Post-indexed and unscaled addressing is rewritten as offsets from each piece's base.
  • The low halves of v8–v15 are kept in x7–x14 (umov/ins) rather than on the stack, so the function needs no stack.
• Callers:
  • nttAtWithW now calls vg_mlkem16_ntt, where it used to inline Readonly.fwdW.
  • So every 16-bit forward NTT of the AArch64 KEM tops (keygen, encaps, and the re-encryption in decaps) runs the new code.
  • The tops' constants gain the table.
• Deleted: Readonly.fwdW, which nothing calls now, with its row bodies (pass1BodyW, pass2BodyW), its literal and its proofs (W.fwdW_ok, fwdContractW, NttR.fwdW_run, checkW_ok). What invW shares with it stays (InW, ctxW, initW_rel, polyIsW_of_outW).
• Decaps's decryption still decodes u into u32s and transforms it with the u32 NTT. Moving it to 16 bits is a follow-up.

Proofs

All by the machine of Proof/MlKem/AArch64/NttR/:

• RunF.checkF_ok: the kernel's run of the whole function (lit_decide, about 2.9 s), from the table as constants. LitF materializes the code and the table.
  • The machine gains the by-element multiplies (Sym, Step: eval_mulHE), ins .d2 and umov .x.
• FwdF.fwdF_run: the run leaves NTT of the input, reduced, in the polynomial.
• AbiF.fwdFast_abi: the ABI.
  • save copies the low halves of v8–v15 to x7–x14 (umovOf_ok).
  • The transform writes none of x7–x14 (core_keeps).
  • restore puts them back (insOf_ok).
  • Nothing writes x19–x28, x30 or sp (fwdFast_keeps).
• VerifiedF:
  • constant time, by taint_decide;
  • Verified against ntt16Contract (abi.withConsts tables), with a satisfying state.
• CallF.ntt16_call: the call, by WP.callV.
  • W.Call.nttAtWithW_ok keeps its statement except for the table it needs (FastReady), so the noise proofs change only their table hypotheses.
  • Optimized.TransformTables and Tables carry the new table (FastTable).

Benchmarks

Gate: keygen ≥2%; encaps and decaps any gain; nothing slower than 0.7%.

Run 39664565308 on N2 (arches=aarch64, cpu_features=-, modules=mlkem768 mlkem1024), against 4a71169, the branch's merge base with main before the restack. Fastest of 8 copies per side; spread ≤ 0.2%. The restack onto main changed none of this PR's code.

Benchmark	Base	Head	Change
mlkem768_keygen	12.95 µs	12.51 µs	3.4% faster
mlkem768_encaps	13.22 µs	13.00 µs	1.7% faster
mlkem768_encaps_from_bytes	13.43 µs	13.21 µs	1.6% faster
mlkem768_decaps	15.19 µs	15.00 µs	1.3% faster
mlkem1024_keygen	20.39 µs	19.82 µs	2.8% faster
mlkem1024_encaps	20.27 µs	19.99 µs	1.4% faster
mlkem1024_encaps_from_bytes	20.54 µs	20.26 µs	1.4% faster
mlkem1024_decaps	23.00 µs	22.70 µs	1.3% faster

The gate is met.

Checks

• Full lake build and Emit.lean --check.
  • src/asm/aarch64/ is regenerated: mlkem16.rs gains vg_mlkem16_ntt and consts.rs gains the table.
  • mlkem768.rs and mlkem1024.rs call the function where they used to inline the old NTT, about 13,000 fewer lines net.
• cargo test for aarch64-unknown-linux-gnu under qemu-aarch64 -cpu max, with Wycheproof, both by default and with VG_CPU_FEATURES=none: all ML-KEM tests pass.
• The Python checks:
  • check_lean_imports, check_lean_modules and check_lean_public_imports;
  • check_registrations and check_lean_speed;
  • check_vectors and check_arch_gates;
  • check_variants and check_mcdt.

🤖 Generated with Claude Code

https://claude.ai/code/session_01FoYWDDEDFq1DFev5X7n2aq

[alex]

N2 benchmarks (run 39664565308, arches=aarch64, cpu_features=-, modules=mlkem768 mlkem1024, against 4a71169, this branch's merge base with main; fastest of 8 copies per side, spread ≤ 0.2%):

Benchmark	Base	Head	Change
mlkem768_keygen	12.95 µs	12.51 µs	3.4% faster
mlkem768_encaps	13.22 µs	13.00 µs	1.7% faster
mlkem768_encaps_from_bytes	13.43 µs	13.21 µs	1.6% faster
mlkem768_decaps	15.19 µs	15.00 µs	1.3% faster
mlkem1024_keygen	20.39 µs	19.82 µs	2.8% faster
mlkem1024_encaps	20.27 µs	19.99 µs	1.4% faster
mlkem1024_encaps_from_bytes	20.54 µs	20.26 µs	1.4% faster
mlkem1024_decaps	23.00 µs	22.70 µs	1.3% faster

This meets the gate: keygen at least 2% faster, encaps and decaps faster, and nothing slower. Next on this PR: delete the forward NTT code this leaves unused (Readonly.fwdW and its proofs). Then I'll restack once #3020 and #3042 merge.

---

Generated by Claude Code

[alex]

src/asm/aarch64/ at this head (7040c17) is not byte-identical to the head the gate run in the description measured (f50412a). The restack onto main brought in #3031's ML-KEM rejection-sampling change and the P-384 work. So I re-ran the gate on this head:

Run 39685056275 on N2 (aarch64, -, mlkem768 mlkem1024), against main at b524fb2. Fastest of 8 copies per side; spread ≤ 0.3%.

Benchmark	Base	Head	Change
mlkem768_keygen	12.73 µs	12.28 µs	3.5% faster
mlkem768_encaps	12.98 µs	12.79 µs	1.5% faster
mlkem768_encaps_from_bytes	13.19 µs	12.99 µs	1.5% faster
mlkem768_decaps	14.97 µs	14.76 µs	1.4% faster
mlkem1024_keygen	20.00 µs	19.42 µs	2.9% faster
mlkem1024_encaps	19.87 µs	19.58 µs	1.4% faster
mlkem1024_encaps_from_bytes	20.14 µs	19.85 µs	1.4% faster
mlkem1024_decaps	22.59 µs	22.31 µs	1.2% faster

The gate still holds on this head: keygen is at least 2% faster, encaps and decaps are faster, and nothing is slower. CI is green.

---

Generated by Claude Code