Skip to content

ML-KEM AArch64: 16-bit paired key generation, encapsulation and decapsulation - #2849

Merged
reaperhulk merged 39 commits into
mainfrom
claude/a64-mlkem16
Oct 19, 2026
Merged

ML-KEM AArch64: 16-bit paired key generation, encapsulation and decapsulation#2849
reaperhulk merged 39 commits into
mainfrom
claude/a64-mlkem16

Conversation

@alex

@alex alex commented Oct 19, 2026

Copy link
Copy Markdown
Member

Part of #991.

On AArch64, ML-KEM-768 and ML-KEM-1024 key generation, encapsulation and decapsulation now keep polynomials as u16 coefficients all the way through. Each polynomial sits in the first 512 bytes of its 1024-byte slot. The old paired whole-KEM compositions kept u32 coefficients, and these replace them.

Changes

  • Impl: Impl/MlKem/AArch64/Kem16.lean, plus the remaining 16-bit pieces it uses (widen, a 16-bit CBD call in the shared noise sampler). The new pieces:
    • addW
    • dotW
    • nttInvAtWithW
    • the 12-bit decW/encW codecs
  • Proofs: Proof/MlKem/AArch64/Kem16/* proves all six artifacts Verified (correct, constant time, stack depth).
    • The KB invariants are generalized over width (BInv := BInvW false, …).
    • half_sub carries frames over the 512-byte halves to the 1024-byte slots.
  • Registration: Generic/KeccakX2/Keccak/AArch64/MlKem.lean now binds the 16-bit keyGen/encaps/decaps. src/asm/aarch64/ is regenerated, and the new module mlkem16.rs holds vg_mlkem16_cbd2.
  • Deletions: everything the u32 paired compositions leave dead, deleted at the declaration level, as ML-KEM x86-64: fused encapsulation with 16-bit polynomials; delete the u32 code with AVX2 #2830 did:
    • Optimized: Kem*, KeyGen*, DecapsTiming, …
    • FixedHash: KeyGenFinal/Verified/Using, DecapsVerified/Correct, …
    • the u32 vRest/encryptRest/encaps compositions and their abbreviations
    • narrow
    • Modules left empty are removed, and their importers are spliced. Every Api is still registered, so no follow-up Spec PR is needed.

No changes to TCB/ or Spec/.

Speed

Benchmarks run 39186145283, Neoverse N2, fastest of 8 copies, all features:

op ML-KEM-768 ML-KEM-1024
keygen −4.5% −4.7%
encaps −5.6% −5.6%
encaps_from_bytes −5.5% −5.4%
decaps −4.4% −4.3%

neon and none are 0.0%: that code is unchanged. No row is slower.

Checks

These pass locally:

  • full lake build
  • Emit.lean --check
  • check_lean_imports, check_lean_modules, check_registrations, check_lean_speed, check_vectors, check_arch_gates, check_variants, check_mcdt
  • algorithms_table.py (table unchanged)
  • cargo fmt --check and clippy -D warnings
  • the aarch64 ML-KEM ACVP and Wycheproof tests under qemu

Follow-ups

  • The w16 parameters on the shared noise sampler are now always true, so they can be removed.
  • subW/hsubOp were already there before this PR and are left as they are.
  • Native 16-bit codecs for the other widths.

Generated by Claude Code

claude added 30 commits October 18, 2026 14:09
@claude
ML-KEM AArch64: vg_mlkem16_cbd2, SamplePolyCBD2 with 16-bit coefficients
901e8a4 
Two TBLs look up the low and high bytes of (f - 2) mod q for each nibble
f <= 4, and ZIP1/ZIP2 interleave them into halfwords: 14 vector
instructions per 32 coefficients after the nibble arithmetic, against 46
for the u32 version's widening and reduction.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64: the NTT's arithmetic as lemmas of its own (p1_arith, …
c3eb79e 
…p2_arith)

pass1Body and pass2Body name their arithmetic (pass1Arith, pass2Arith) and
pass 2's loads, so that a pass with other loads or stores can share them;
the instructions are the same, and so is the emitted code.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64: memory lemmas of 16-bit coefficients (Coeff16), used …
9fb0d2c 
…by vg_mlkem16_cbd2

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64: the 16-bit NTT of a polynomial stored with 16 bits (f…
7203be8 
…wdW)

Each row of pass 1 is one load (no uzp1 from u32s), and each vector of
pass 2 one store (no widening zips): 160 fewer instructions per
transform. The arithmetic is pass1Arith/pass2Arith, shared with the u32
version; wp_stores takes the region's length (wp_storesN), and pass 2's
loads come from where they are in memory (p2_loadsOf).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64: the inverse NTT's arithmetic as lemmas of its own (ip…
6d30c3d 
…1_arith, ip2_arith)

invPass1Body and invPass2Body name their arithmetic (invPass1Arith,
invPass2Arith); the instructions, and the emitted code, are the same.
invPass1BodyW and invPass2BodyW load and store 16-bit coefficients.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64: the 16-bit inverse NTT of a polynomial stored with 16…
f01789d 
… bits (invW)

Each vector of pass 1 is one load and each row of pass 2 one store (no
uzp1 from u32s, no widening zips). CoeffsInW, coeffAtW16_write: 16-bit
coefficients written eight at a time at any offset.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64: the fused dot product of 16-bit polynomials (dotW) an…
c323e99 
…d its lane lemmas

Each iteration takes 8 pairs of each term: UZP1/UZP2 of halfwords and
UMLAL/UMLAL2 into word accumulators. The proof of the loop follows.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM: 16-bit coefficients from their bytes in a target-independent …
d19fcdf 
…module (Mem16Bytes)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64 dotW: one term's gather and accumulation (gatherW_ok, …
d26269e 
…accumulateW_ok)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64 dotW: a whole term, all the terms, and the finish with…
83dc2e2 
… any registers

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64 dotW: the store of sixteen coefficients (storeW_ok)
639793d 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64 dotW: an iteration's sums, the step and the loop
f053a48 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64: dotW writes the fused dot product with 16-bit coeffic…
66b5af4 
…ients (dotW_ok)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64: ByteEncode12 of a polynomial stored with 16 bits (enc…
37f9b35 
…W_ok)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64: ByteDecode12 to a polynomial stored with 16 bits (dec…
680d95b 
…W_ok)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64 dotW: base registers for the output and for each opera…
f7614bf 
…nd of the terms

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64: sums and differences of polynomials stored with 16 bi…
f27e993 
…ts (addW_ok, subW_ok)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64 sampler: a w16 option that stores coefficients 16 bits…
e49c202 
… apart (default off, same code)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM: a 32-bit store at a coefficient stored with 16 or 32 bits (cv…
df8a4af 
…al_writeW)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64 SampleNTT loop: invariants over a coefficient width (1…
b876476 
…6 or 32 bits)

The 32-bit names are the w16 = false instances; the code is unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64 rejection sampler: vector batches, loop and paired pie…
6a7721e 
…ces over a coefficient width

The 16-bit batch stores the compacted halfwords themselves; the 32-bit names
are the w16 = false instances, and the code is unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64 paired sampler: initial, resume and timing pieces over…
845cbf3 
… a coefficient width

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64 paired sampler: extra blocks, finish, pair and their t…
c83b9b2 
…iming over a coefficient width

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64 paired sampler: the last sample and its fallback over …
5fe71fd 
…a coefficient width

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64: matrix invariants and the paired and hashing matrices…
43c7505 
… over a coefficient width

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64: the encryption invariant over a coefficient width
6f3bc2e 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
Merge remote-tracking branch 'origin/main' into claude/a64-mlkem16
4daa633
@claude
ML-KEM AArch64: the coefficient width from Mem16 (cw, cAddr), shared …
e6e59a6 
…with x86-64

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64 noise: CBD and the NTT over a coefficient width (vg_ml…
b22d125 
…kem16_cbd2 and the 16-bit-stored transforms inlined)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64: widen and narrow a polynomial between 16-bit and 32-b…
4974bc1 
…it coefficients

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
claude added 9 commits October 19, 2026 00:55
@claude
ML-KEM AArch64: 16-bit encryption and key generation code, and its bu…
39bf1fd 
…ilding blocks in the KEM context

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64: 16-bit noise and ciphertext rows of u
2ffebc7 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64: 16-bit ciphertext component v
36088c6 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64: the 16-bit hashing encapsulation is correct
594f312 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64: the 16-bit encapsulations are constant time and verified
c364d0e 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64: the 16-bit decapsulations are correct, constant time …
2133fc4 
…and verified

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64: the 16-bit key generations are correct, constant time…
32d0c2f 
… and verified

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64: register the 16-bit paired key generation, encapsulat…
3e419fa 
…ion and decapsulation

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@claude
ML-KEM AArch64: delete the u32 paired code nothing emits any more
17363b9 
The 16-bit key generation, encapsulation and decapsulation replace the
u32 paired compositions, so their implementations and proofs, and the
lemmas only they used, are deleted.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ma2fA9rUTPXEEQjYyCMCRw
@reaperhulk
reaperhulk added this pull request to the merge queue Oct 19, 2026
Merged via the queue into main with commit 416f9df Oct 19, 2026
14 checks passed
@reaperhulk
reaperhulk deleted the claude/a64-mlkem16 branch October 19, 2026 04:44
@alex alex mentioned this pull request Oct 19, 2026
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@alex @reaperhulk @claude